feat: integrate AI summary into SessionContextService (P0 phase 1)

- SessionContextService 新增 ai_summary_service 可选字段
- 新增 new_with_ai_summary() 构造函数
- create_summary() 改为异步方法,优先使用 AI 摘要
- AI 摘要失败时自动降级到本地关键词提取
- 将 get_effective_context/preheat_session_context 改为异步
- 修复所有相关测试为 tokio::test
- 全部 11 个测试通过(7 session_context + 4 ai_summary)
This commit is contained in:
coso
2026-03-12 19:05:28 +08:00
parent 75309bd2ef
commit 621ab3ee73
2 changed files with 437 additions and 56 deletions
@@ -4,7 +4,6 @@
use proxycast_core::general_chat::{ChatMessage, MessageRole};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
/// AI 摘要请求
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -2,13 +2,15 @@
//!
//! 提供会话上下文的持久化、恢复和智能管理功能,解决 AI 对话中的上下文丢失问题
use crate::general_chat::{ChatMessage, MessageRole};
use proxycast_core::database::dao::general_chat::GeneralChatDao;
use crate::ai_summary_service::AISummaryService;
use proxycast_core::database::dao::chat::{ChatDao, ChatMessage as UnifiedChatMessage, ChatMode};
use proxycast_core::database::load_pending_general_session_messages;
use proxycast_core::general_chat::{ChatMessage, MessageRole};
use rusqlite::Connection;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use tracing::{debug, info};
use tracing::{debug, info, warn};
/// 会话上下文摘要
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -61,6 +63,8 @@ pub struct SessionContextService {
config: ContextWindowConfig,
/// 会话摘要缓存
summary_cache: Arc<Mutex<HashMap<String, SessionSummary>>>,
/// AI 摘要服务(可选)
ai_summary_service: Option<Arc<AISummaryService>>,
}
impl SessionContextService {
@@ -70,18 +74,34 @@ impl SessionContextService {
db_connection,
config,
summary_cache: Arc::new(Mutex::new(HashMap::new())),
ai_summary_service: None,
}
}
/// 创建带 AI 摘要服务的会话上下文服务
pub fn new_with_ai_summary(
db_connection: Arc<Mutex<Connection>>,
config: ContextWindowConfig,
ai_summary_service: Arc<AISummaryService>,
) -> Self {
Self {
db_connection,
config,
summary_cache: Arc::new(Mutex::new(HashMap::new())),
ai_summary_service: Some(ai_summary_service),
}
}
/// 获取会话的有效上下文
///
/// 返回适合发送给 AI 的消息列表,包括摘要和最近的消息
pub fn get_effective_context(&self, session_id: &str) -> Result<Vec<ChatMessage>, String> {
pub async fn get_effective_context(
&self,
session_id: &str,
) -> Result<Vec<ChatMessage>, String> {
let conn = self.db_connection.lock().map_err(|e| e.to_string())?;
// 获取所有消息
let all_messages = GeneralChatDao::get_messages(&conn, session_id, None, None)
.map_err(|e| format!("获取消息失败: {e}"))?;
let all_messages = Self::load_session_messages(&conn, session_id)?;
if all_messages.is_empty() {
return Ok(vec![]);
@@ -97,10 +117,11 @@ impl SessionContextService {
// 需要应用上下文管理
self.apply_context_window_management(session_id, all_messages)
.await
}
/// 应用上下文窗口管理
fn apply_context_window_management(
async fn apply_context_window_management(
&self,
session_id: &str,
all_messages: Vec<ChatMessage>,
@@ -110,7 +131,7 @@ impl SessionContextService {
// 如果启用智能摘要且消息数量超过阈值
if self.config.enable_smart_summary && all_messages.len() > self.config.summary_threshold {
// 尝试获取或创建摘要
if let Ok(summary) = self.get_or_create_summary(session_id, &all_messages) {
if let Ok(summary) = self.get_or_create_summary(session_id, &all_messages).await {
// 添加摘要作为系统消息
let summary_message = ChatMessage {
id: format!("summary-{session_id}"),
@@ -156,6 +177,148 @@ impl SessionContextService {
Ok(result)
}
fn load_session_messages(
conn: &Connection,
session_id: &str,
) -> Result<Vec<ChatMessage>, String> {
let session =
ChatDao::get_session(conn, session_id).map_err(|e| format!("获取统一会话失败: {e}"))?;
if let Some(session) = session {
if session.mode != ChatMode::General {
debug!(
"会话 {} 不是 general 模式,跳过上下文加载: {:?}",
session_id, session.mode
);
return Ok(vec![]);
}
let unified_messages = ChatDao::get_messages(conn, session_id, None)
.map_err(|e| format!("获取统一消息失败: {e}"))?;
return Ok(unified_messages
.into_iter()
.map(Self::convert_unified_message)
.collect());
}
debug!(
"会话 {} 未命中 unified chat,回退待迁移 general 历史表",
session_id
);
Self::load_pending_general_messages(conn, session_id)
}
fn load_pending_general_messages(
conn: &Connection,
session_id: &str,
) -> Result<Vec<ChatMessage>, String> {
let messages = load_pending_general_session_messages(conn, session_id)
.map_err(|e| format!("查询待迁移 general 消息失败: {e}"))?;
Ok(messages
.into_iter()
.map(|message| ChatMessage {
id: message.id,
session_id: message.session_id,
role: Self::convert_legacy_role(&message.role),
content: message.content,
blocks: None,
status: "complete".to_string(),
created_at: message.created_at,
metadata: None,
})
.collect())
}
fn convert_unified_message(message: UnifiedChatMessage) -> ChatMessage {
ChatMessage {
id: message.id.to_string(),
session_id: message.session_id,
role: Self::convert_unified_role(&message.role),
content: Self::extract_unified_text_content(&message.content),
blocks: None,
status: "complete".to_string(),
created_at: Self::parse_unified_timestamp(&message.created_at),
metadata: message.metadata,
}
}
fn convert_unified_role(role: &str) -> MessageRole {
match role {
"assistant" => MessageRole::Assistant,
"system" => MessageRole::System,
_ => MessageRole::User,
}
}
fn convert_legacy_role(role: &str) -> MessageRole {
match role {
"assistant" => MessageRole::Assistant,
"system" => MessageRole::System,
_ => MessageRole::User,
}
}
fn extract_unified_text_content(content: &serde_json::Value) -> String {
match content {
serde_json::Value::Null => String::new(),
serde_json::Value::String(text) => text.clone(),
serde_json::Value::Array(items) => {
let parts: Vec<String> = items
.iter()
.filter_map(Self::extract_unified_text_fragment)
.filter(|text| !text.trim().is_empty())
.collect();
if parts.is_empty() {
content.to_string()
} else {
parts.join("\n")
}
}
serde_json::Value::Object(_) => {
Self::extract_unified_text_fragment(content).unwrap_or_else(|| content.to_string())
}
_ => content.to_string(),
}
}
fn extract_unified_text_fragment(value: &serde_json::Value) -> Option<String> {
match value {
serde_json::Value::Null => None,
serde_json::Value::String(text) => Some(text.clone()),
serde_json::Value::Array(items) => {
let parts: Vec<String> = items
.iter()
.filter_map(Self::extract_unified_text_fragment)
.filter(|text| !text.trim().is_empty())
.collect();
if parts.is_empty() {
None
} else {
Some(parts.join("\n"))
}
}
serde_json::Value::Object(map) => map
.get("text")
.and_then(|value| value.as_str())
.or_else(|| map.get("content").and_then(|value| value.as_str()))
.or_else(|| map.get("value").and_then(|value| value.as_str()))
.map(ToString::to_string),
_ => None,
}
}
fn parse_unified_timestamp(timestamp: &str) -> i64 {
chrono::DateTime::parse_from_rfc3339(timestamp)
.map(|value| value.timestamp_millis())
.or_else(|_| timestamp.parse::<i64>())
.unwrap_or_default()
}
/// 选择最近的消息,确保不超过配置限制
fn select_recent_messages(
&self,
@@ -185,7 +348,7 @@ impl SessionContextService {
}
/// 获取或创建会话摘要
fn get_or_create_summary(
async fn get_or_create_summary(
&self,
session_id: &str,
messages: &[ChatMessage],
@@ -199,7 +362,7 @@ impl SessionContextService {
}
// 创建新摘要
let summary = self.create_summary(session_id, messages)?;
let summary = self.create_summary(session_id, messages).await?;
// 缓存摘要
{
@@ -210,8 +373,8 @@ impl SessionContextService {
Ok(summary)
}
/// 创建会话摘要
fn create_summary(
/// 创建会话摘要(优先使用 AI,失败时降级到本地)
async fn create_summary(
&self,
session_id: &str,
messages: &[ChatMessage],
@@ -224,6 +387,39 @@ impl SessionContextService {
let summary_count = (messages.len() * 2 / 3).min(self.config.summary_threshold);
let messages_to_summarize = &messages[..summary_count];
// 尝试使用 AI 摘要
if let Some(ai_service) = &self.ai_summary_service {
match ai_service.generate_summary(messages_to_summarize).await {
Ok(ai_summary) => {
info!("AI 摘要生成成功,会话: {}", session_id);
let last_message = messages_to_summarize.last().unwrap();
return Ok(SessionSummary {
session_id: session_id.to_string(),
summary: ai_summary.summary,
key_topics: ai_summary.key_topics,
decisions: ai_summary.decisions,
created_at: chrono::Utc::now().timestamp_millis(),
message_count: summary_count as i32,
last_message_id: last_message.id.clone(),
});
}
Err(e) => {
warn!("AI 摘要生成失败,降级到本地摘要: {}", e);
}
}
}
// 降级到本地关键词提取
self.create_summary_local(session_id, messages_to_summarize)
}
/// 本地关键词提取摘要(保留作为降级方案)
fn create_summary_local(
&self,
session_id: &str,
messages_to_summarize: &[ChatMessage],
) -> Result<SessionSummary, String> {
// 提取关键信息
let mut key_topics = Vec::new();
let mut decisions = Vec::new();
@@ -252,7 +448,7 @@ impl SessionContextService {
let summary_content = if content_parts.len() > 10 {
format!(
"本会话包含 {} 条消息,主要讨论了以下内容:\n{}",
summary_count,
messages_to_summarize.len(),
content_parts
.join("\n")
.chars()
@@ -271,7 +467,7 @@ impl SessionContextService {
key_topics,
decisions,
created_at: chrono::Utc::now().timestamp_millis(),
message_count: summary_count as i32,
message_count: messages_to_summarize.len() as i32,
last_message_id: last_message.id.clone(),
})
}
@@ -340,11 +536,8 @@ impl SessionContextService {
pub fn get_session_stats(&self, session_id: &str) -> Result<SessionStats, String> {
let conn = self.db_connection.lock().map_err(|e| e.to_string())?;
let message_count = GeneralChatDao::get_message_count(&conn, session_id)
.map_err(|e| format!("获取消息数量失败: {e}"))? as usize;
let messages = GeneralChatDao::get_messages(&conn, session_id, None, None)
.map_err(|e| format!("获取消息失败: {e}"))?;
let messages = Self::load_session_messages(&conn, session_id)?;
let message_count = messages.len();
let total_characters: usize = messages.iter().map(|m| m.content.len()).sum();
let user_messages = messages
@@ -378,15 +571,14 @@ impl SessionContextService {
}
/// 预热会话上下文(提前创建摘要)
pub fn preheat_session_context(&self, session_id: &str) -> Result<(), String> {
pub async fn preheat_session_context(&self, session_id: &str) -> Result<(), String> {
let conn = self.db_connection.lock().map_err(|e| e.to_string())?;
let messages = GeneralChatDao::get_messages(&conn, session_id, None, None)
.map_err(|e| format!("获取消息失败: {e}"))?;
let messages = Self::load_session_messages(&conn, session_id)?;
if messages.len() > self.config.summary_threshold {
debug!("为会话 {} 预热上下文", session_id);
self.get_or_create_summary(session_id, &messages)?;
self.get_or_create_summary(session_id, &messages).await?;
info!("会话 {} 上下文预热完成", session_id);
}
@@ -410,12 +602,54 @@ pub struct SessionStats {
#[cfg(test)]
mod tests {
use super::*;
use crate::general_chat::ChatSession;
use proxycast_core::database::dao::chat::{
ChatDao, ChatMessage as UnifiedChatMessage, ChatMode, ChatSession as UnifiedChatSession,
};
use proxycast_core::general_chat::ChatSession;
use rusqlite::Connection;
fn setup_test_db() -> Connection {
let conn = Connection::open_in_memory().unwrap();
conn.execute(
"CREATE TABLE agent_sessions (
id TEXT PRIMARY KEY,
model TEXT NOT NULL,
system_prompt TEXT,
title TEXT,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL,
working_dir TEXT,
execution_strategy TEXT NOT NULL DEFAULT 'react'
)",
[],
)
.unwrap();
conn.execute(
"CREATE TABLE agent_messages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
session_id TEXT NOT NULL,
role TEXT NOT NULL,
content_json TEXT NOT NULL,
timestamp TEXT NOT NULL,
tool_calls_json TEXT,
tool_call_id TEXT,
FOREIGN KEY (session_id) REFERENCES agent_sessions(id) ON DELETE CASCADE
)",
[],
)
.unwrap();
conn.execute(
"CREATE TABLE settings (
key TEXT PRIMARY KEY,
value TEXT NOT NULL
)",
[],
)
.unwrap();
// 创建会话表
conn.execute(
"CREATE TABLE general_chat_sessions (
@@ -477,8 +711,86 @@ mod tests {
messages
}
#[test]
fn test_context_service_creation() {
fn create_unified_general_session(session_id: &str) -> UnifiedChatSession {
let now = chrono::Utc::now().to_rfc3339();
UnifiedChatSession {
id: session_id.to_string(),
mode: ChatMode::General,
title: Some("测试会话".to_string()),
system_prompt: None,
model: None,
provider_type: None,
credential_uuid: None,
metadata: None,
created_at: now.clone(),
updated_at: now,
}
}
fn create_unified_messages(session_id: &str, count: usize) -> Vec<UnifiedChatMessage> {
let base_time = chrono::Utc::now();
(0..count)
.map(|index| {
let role = if index % 2 == 0 { "user" } else { "assistant" };
let content = format!("这是第 {} 条消息,包含一些测试内容", index + 1);
UnifiedChatMessage {
id: 0,
session_id: session_id.to_string(),
role: role.to_string(),
content: serde_json::json!([{ "type": "text", "text": content }]),
tool_calls: None,
tool_call_id: None,
metadata: None,
created_at: (base_time + chrono::Duration::milliseconds(index as i64))
.to_rfc3339(),
}
})
.collect()
}
fn insert_legacy_session(conn: &Connection, session: &ChatSession) {
conn.execute(
"INSERT INTO general_chat_sessions (id, name, created_at, updated_at, metadata)
VALUES (?1, ?2, ?3, ?4, ?5)",
rusqlite::params![
session.id,
session.name,
session.created_at,
session.updated_at,
Option::<String>::None,
],
)
.unwrap();
}
fn insert_legacy_message(conn: &Connection, message: &ChatMessage) {
let role = match message.role {
MessageRole::User => "user",
MessageRole::Assistant => "assistant",
MessageRole::System => "system",
};
conn.execute(
"INSERT INTO general_chat_messages (id, session_id, role, content, blocks, status, created_at, metadata)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
rusqlite::params![
message.id,
message.session_id,
role,
message.content,
Option::<String>::None,
message.status,
message.created_at,
Option::<String>::None,
],
)
.unwrap();
}
#[tokio::test]
async fn test_context_service_creation() {
let conn = Arc::new(Mutex::new(setup_test_db()));
let config = ContextWindowConfig::default();
let service = SessionContextService::new(conn, config);
@@ -486,8 +798,8 @@ mod tests {
assert!(service.summary_cache.lock().unwrap().is_empty());
}
#[test]
fn test_get_effective_context_small_session() {
#[tokio::test]
async fn test_get_effective_context_small_session() {
let conn = Arc::new(Mutex::new(setup_test_db()));
let config = ContextWindowConfig {
max_messages: 50,
@@ -497,28 +809,19 @@ mod tests {
};
let service = SessionContextService::new(conn.clone(), config);
// 创建会话
let session = ChatSession {
id: "test-session".to_string(),
name: "测试会话".to_string(),
created_at: chrono::Utc::now().timestamp_millis(),
updated_at: chrono::Utc::now().timestamp_millis(),
metadata: None,
};
{
let conn_guard = conn.lock().unwrap();
GeneralChatDao::create_session(&conn_guard, &session).unwrap();
ChatDao::create_session(&conn_guard, &create_unified_general_session("test-session"))
.unwrap();
// 添加少量消息
let messages = create_test_messages("test-session", 5);
let messages = create_unified_messages("test-session", 5);
for msg in &messages {
GeneralChatDao::add_message(&conn_guard, msg).unwrap();
ChatDao::add_message(&conn_guard, msg).unwrap();
}
}
// 获取有效上下文
let context = service.get_effective_context("test-session").unwrap();
let context = service.get_effective_context("test-session").await.unwrap();
assert_eq!(context.len(), 5);
}
@@ -528,22 +831,14 @@ mod tests {
let config = ContextWindowConfig::default();
let service = SessionContextService::new(conn.clone(), config);
// 创建会话和消息
let session = ChatSession {
id: "test-session".to_string(),
name: "测试会话".to_string(),
created_at: chrono::Utc::now().timestamp_millis(),
updated_at: chrono::Utc::now().timestamp_millis(),
metadata: None,
};
{
let conn_guard = conn.lock().unwrap();
GeneralChatDao::create_session(&conn_guard, &session).unwrap();
ChatDao::create_session(&conn_guard, &create_unified_general_session("test-session"))
.unwrap();
let messages = create_test_messages("test-session", 10);
let messages = create_unified_messages("test-session", 10);
for msg in &messages {
GeneralChatDao::add_message(&conn_guard, msg).unwrap();
ChatDao::add_message(&conn_guard, msg).unwrap();
}
}
@@ -584,4 +879,91 @@ mod tests {
service.clear_summary_cache(Some("session-1")).unwrap();
assert_eq!(service.summary_cache.lock().unwrap().len(), 0);
}
#[tokio::test]
async fn test_get_effective_context_falls_back_to_legacy_messages() {
let conn = Arc::new(Mutex::new(setup_test_db()));
let service = SessionContextService::new(conn.clone(), ContextWindowConfig::default());
let legacy_session = ChatSession {
id: "legacy-session".to_string(),
name: "旧会话".to_string(),
created_at: chrono::Utc::now().timestamp_millis(),
updated_at: chrono::Utc::now().timestamp_millis(),
metadata: None,
};
{
let conn_guard = conn.lock().unwrap();
insert_legacy_session(&conn_guard, &legacy_session);
for message in create_test_messages("legacy-session", 3) {
insert_legacy_message(&conn_guard, &message);
}
}
let context = service.get_effective_context("legacy-session").await.unwrap();
assert_eq!(context.len(), 3);
assert_eq!(context[0].content, "这是第 1 条消息,包含一些测试内容");
}
#[tokio::test]
async fn test_get_effective_context_skips_legacy_fallback_after_general_migration_completed() {
let conn = Arc::new(Mutex::new(setup_test_db()));
let service = SessionContextService::new(conn.clone(), ContextWindowConfig::default());
let legacy_session = ChatSession {
id: "legacy-session".to_string(),
name: "旧会话".to_string(),
created_at: chrono::Utc::now().timestamp_millis(),
updated_at: chrono::Utc::now().timestamp_millis(),
metadata: None,
};
{
let conn_guard = conn.lock().unwrap();
conn_guard
.execute(
"INSERT INTO settings (key, value) VALUES (?1, ?2)",
rusqlite::params!["migrated_general_chat_to_unified", "true"],
)
.unwrap();
insert_legacy_session(&conn_guard, &legacy_session);
for message in create_test_messages("legacy-session", 3) {
insert_legacy_message(&conn_guard, &message);
}
}
let context = service.get_effective_context("legacy-session").await.unwrap();
assert!(context.is_empty());
}
#[tokio::test]
async fn test_non_general_unified_session_returns_empty_context() {
let conn = Arc::new(Mutex::new(setup_test_db()));
let service = SessionContextService::new(conn.clone(), ContextWindowConfig::default());
{
let conn_guard = conn.lock().unwrap();
let now = chrono::Utc::now().to_rfc3339();
ChatDao::create_session(
&conn_guard,
&UnifiedChatSession {
id: "agent-session".to_string(),
mode: ChatMode::Agent,
title: Some("Agent 会话".to_string()),
system_prompt: None,
model: Some("claude-sonnet-4".to_string()),
provider_type: None,
credential_uuid: None,
metadata: None,
created_at: now.clone(),
updated_at: now,
},
)
.unwrap();
}
let context = service.get_effective_context("agent-session").await.unwrap();
assert!(context.is_empty());
}
}