gardel/Mofox-Core
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #5 from MoFox-Studio/dev

Browse Source 
Dev
This commit is contained in:
yishan
2025-09-23 12:34:09 +08:00
committed by GitHub
parent 0c70749a66 a02c1f095c
commit f4e18e2657
143 changed files with 2566 additions and 4247 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
README.md
View File

@@ -29,7 +29,7 @@
## 📖 项目介绍
**MoFox_Bot** 是一个基于 [MaiCore](https://github.com/MaiM-with-u/MaiBot) `0.10.0 snapshot.5` 版本的增强型 `fork` 项目。
我们在保留原版所有功能的基础上，进行了一系列的改进和功能拓展，致力于提供更强的稳定性、更丰富的功能和更流畅的用户体验
我们在保留原版几乎所有功能的基础上，进行了一系列的改进和功能拓展，致力于提供更强的稳定性、更丰富的功能和更流畅的用户体验
> [!IMPORTANT]
> **第三方项目声明**

6
bot.py
View File

@@ -193,9 +193,11 @@ class MaiBotMain(BaseMain):
logger.error(f"数据库连接初始化失败: {e}")
raise e
async def initialize_database_async(self):
"""异步初始化数据库表结构"""
logger.info("正在初始化数据库表结构...")
try:
init_db()
await init_db()
logger.info("数据库表结构初始化完成")
except Exception as e:
logger.error(f"数据库表结构初始化失败: {e}")
@@ -229,6 +231,8 @@ if __name__ == "__main__":
try:
# 执行初始化和任务调度
loop.run_until_complete(main_system.initialize())
# 异步初始化数据库表结构
loop.run_until_complete(maibot.initialize_database_async())
initialize_lpmm_knowledge()
# Schedule tasks returns a future that runs forever.
# We can run console_input_loop concurrently.

3
pyproject.toml
View File

@@ -72,6 +72,9 @@ dependencies = [
"uvicorn>=0.35.0",
"watchdog>=6.0.0",
"websockets>=15.0.1",
"aiomysql>=0.2.0",
"aiosqlite>=0.21.0",
"inkfox>=0.1.0",
]
[[tool.uv.index]]

5
requirements.txt
View File

@@ -1,4 +1,6 @@
sqlalchemy
aiosqlite
aiomysql
APScheduler
aiohttp
aiohttp-cors
@@ -67,4 +69,5 @@ google-generativeai
lunar_python
fuzzywuzzy
python-multipart
aiofiles
aiofiles
inkfox

0
rust_image/Cargo.toml Normal file
View File

31
rust_video/Cargo.toml
View File

@@ -1,31 +0,0 @@
[package]
name = "rust_video"
version = "0.1.0"
edition = "2021"
authors = ["VideoAnalysis Team"]
description = "Ultra-fast video keyframe extraction tool in Rust"
license = "MIT"
[dependencies]
anyhow = "1.0"
clap = { version = "4.0", features = ["derive"] }
rayon = "1.11"
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
chrono = { version = "0.4", features = ["serde"] }
# PyO3 dependencies
pyo3 = { version = "0.22", features = ["extension-module"] }
[lib]
name = "rust_video"
crate-type = ["cdylib"]
[profile.release]
opt-level = 3
lto = true
codegen-units = 1
panic = "abort"
strip = true

15
rust_video/pyproject.toml
View File

@@ -1,15 +0,0 @@
[build-system]
requires = ["maturin>=1.9,<2.0"]
build-backend = "maturin"
[project]
name = "rust_video"
requires-python = ">=3.8"
classifiers = [
"Programming Language :: Rust",
"Programming Language :: Python :: Implementation :: CPython",
"Programming Language :: Python :: Implementation :: PyPy",
]
dynamic = ["version"]
[tool.maturin]
features = ["pyo3/extension-module"]

391
rust_video/rust_video.pyi
View File

@@ -1,391 +0,0 @@
"""
Rust Video Keyframe Extractor - Python Type Hints
Ultra-fast video keyframe extraction tool with SIMD optimization.
"""
from typing import Dict, List, Optional, Tuple, Union, Any
from pathlib import Path
class PyVideoFrame:
"""
Python绑定的视频帧结构
表示一个视频帧，包含帧编号、尺寸和像素数据。
"""
frame_number: int
"""帧编号"""
width: int
"""帧宽度（像素）"""
height: int
"""帧高度（像素）"""
def __init__(self, frame_number: int, width: int, height: int, data: List[int]) -> None:
"""
创建新的视频帧
Args:
frame_number: 帧编号
width: 帧宽度
height: 帧高度
data: 像素数据（灰度值列表）
"""
...
def get_data(self) -> List[int]:
"""
获取帧的像素数据
Returns:
像素数据列表（灰度值）
"""
...
def calculate_difference(self, other: 'PyVideoFrame') -> float:
"""
计算与另一帧的差异
Args:
other: 要比较的另一帧
Returns:
帧差异值（0-255范围）
"""
...
def calculate_difference_simd(self, other: 'PyVideoFrame', block_size: Optional[int] = None) -> float:
"""
使用SIMD优化计算帧差异
Args:
other: 要比较的另一帧
block_size: 处理块大小，默认8192
Returns:
帧差异值（0-255范围）
"""
...
class PyPerformanceResult:
"""
性能测试结果
包含详细的性能统计信息。
"""
test_name: str
"""测试名称"""
video_file: str
"""视频文件名"""
total_time_ms: float
"""总处理时间（毫秒）"""
frame_extraction_time_ms: float
"""帧提取时间（毫秒）"""
keyframe_analysis_time_ms: float
"""关键帧分析时间（毫秒）"""
total_frames: int
"""总帧数"""
keyframes_extracted: int
"""提取的关键帧数"""
keyframe_ratio: float
"""关键帧比例（百分比）"""
processing_fps: float
"""处理速度（帧每秒）"""
threshold: float
"""检测阈值"""
optimization_type: str
"""优化类型"""
simd_enabled: bool
"""是否启用SIMD"""
threads_used: int
"""使用的线程数"""
timestamp: str
"""时间戳"""
def to_dict(self) -> Dict[str, Any]:
"""
转换为Python字典
Returns:
包含所有结果字段的字典
"""
...
class VideoKeyframeExtractor:
"""
主要的视频关键帧提取器类
提供完整的视频关键帧提取功能，包括SIMD优化和多线程处理。
"""
def __init__(
self,
ffmpeg_path: str = "ffmpeg",
threads: int = 0,
verbose: bool = False
) -> None:
"""
创建关键帧提取器
Args:
ffmpeg_path: FFmpeg可执行文件路径，默认"ffmpeg"
threads: 线程数，0表示自动检测
verbose: 是否启用详细输出
"""
...
def extract_frames(
self,
video_path: str,
max_frames: Optional[int] = None
) -> Tuple[List[PyVideoFrame], int, int]:
"""
从视频中提取帧
Args:
video_path: 视频文件路径
max_frames: 最大提取帧数，None表示提取所有帧
Returns:
(帧列表, 宽度, 高度)
"""
...
def extract_keyframes(
self,
frames: List[PyVideoFrame],
threshold: float,
use_simd: Optional[bool] = None,
block_size: Optional[int] = None
) -> List[int]:
"""
提取关键帧索引
Args:
frames: 视频帧列表
threshold: 检测阈值
use_simd: 是否使用SIMD优化，默认True
block_size: 处理块大小，默认8192
Returns:
关键帧索引列表
"""
...
def save_keyframes(
self,
video_path: str,
keyframe_indices: List[int],
output_dir: str,
max_save: Optional[int] = None
) -> int:
"""
保存关键帧为图片
Args:
video_path: 原视频文件路径
keyframe_indices: 关键帧索引列表
output_dir: 输出目录
max_save: 最大保存数量，默认50
Returns:
实际保存的关键帧数量
"""
...
def benchmark(
self,
video_path: str,
threshold: float,
test_name: str,
max_frames: Optional[int] = None,
use_simd: Optional[bool] = None,
block_size: Optional[int] = None
) -> PyPerformanceResult:
"""
运行性能测试
Args:
video_path: 视频文件路径
threshold: 检测阈值
test_name: 测试名称
max_frames: 最大处理帧数，默认1000
use_simd: 是否使用SIMD优化，默认True
block_size: 处理块大小，默认8192
Returns:
性能测试结果
"""
...
def process_video(
self,
video_path: str,
output_dir: str,
threshold: Optional[float] = None,
max_frames: Optional[int] = None,
max_save: Optional[int] = None,
use_simd: Optional[bool] = None,
block_size: Optional[int] = None
) -> PyPerformanceResult:
"""
完整的处理流程
执行完整的视频关键帧提取和保存流程。
Args:
video_path: 视频文件路径
output_dir: 输出目录
threshold: 检测阈值，默认2.0
max_frames: 最大处理帧数，0表示处理所有帧
max_save: 最大保存数量，默认50
use_simd: 是否使用SIMD优化，默认True
block_size: 处理块大小，默认8192
Returns:
处理结果
"""
...
def get_cpu_features(self) -> Dict[str, bool]:
"""
获取CPU特性信息
Returns:
CPU特性字典，包含AVX2、SSE2等支持信息
"""
...
def get_thread_count(self) -> int:
"""
获取当前配置的线程数
Returns:
配置的线程数
"""
...
def get_configured_threads(self) -> int:
"""
获取配置的线程数
Returns:
配置的线程数
"""
...
def get_actual_thread_count(self) -> int:
"""
获取实际运行的线程数
Returns:
实际运行的线程数
"""
...
def extract_keyframes_from_video(
video_path: str,
output_dir: str,
threshold: Optional[float] = None,
max_frames: Optional[int] = None,
max_save: Optional[int] = None,
ffmpeg_path: Optional[str] = None,
use_simd: Optional[bool] = None,
threads: Optional[int] = None,
verbose: Optional[bool] = None
) -> PyPerformanceResult:
"""
便捷函数：从视频提取关键帧
这是一个便捷函数，封装了完整的关键帧提取流程。
Args:
video_path: 视频文件路径
output_dir: 输出目录
threshold: 检测阈值，默认2.0
max_frames: 最大处理帧数，0表示处理所有帧
max_save: 最大保存数量，默认50
ffmpeg_path: FFmpeg路径，默认"ffmpeg"
use_simd: 是否使用SIMD优化，默认True
threads: 线程数，0表示自动检测
verbose: 是否启用详细输出，默认False
Returns:
处理结果
Example:
>>> result = extract_keyframes_from_video(
... "video.mp4",
... "./output",
... threshold=2.5,
... max_save=30,
... verbose=True
... )
>>> print(f"提取了 {result.keyframes_extracted} 个关键帧")
"""
...
def get_system_info() -> Dict[str, Any]:
"""
获取系统信息
Returns:
系统信息字典，包含：
- threads: 可用线程数
- avx2_supported: 是否支持AVX2（x86_64）
- sse2_supported: 是否支持SSE2（x86_64）
- simd_supported: 是否支持SIMD（非x86_64）
- version: 库版本
Example:
>>> info = get_system_info()
>>> print(f"线程数: {info['threads']}")
>>> print(f"AVX2支持: {info.get('avx2_supported', False)}")
"""
...
# 类型别名
VideoPath = Union[str, Path]
"""视频文件路径类型"""
OutputPath = Union[str, Path]
"""输出路径类型"""
FrameData = List[int]
"""帧数据类型（像素值列表）"""
KeyframeIndices = List[int]
"""关键帧索引类型"""
# 常量
DEFAULT_THRESHOLD: float = 2.0
"""默认检测阈值"""
DEFAULT_BLOCK_SIZE: int = 8192
"""默认处理块大小"""
DEFAULT_MAX_SAVE: int = 50
"""默认最大保存数量"""
MAX_FRAME_DIFFERENCE: float = 255.0
"""最大帧差异值"""
# 版本信息
__version__: str = "0.1.0"
"""库版本"""

831
rust_video/src/lib.rs
View File

@@ -1,831 +0,0 @@
use pyo3::prelude::*;
use anyhow::{Context, Result};
use chrono::prelude::*;
use rayon::prelude::*;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs;
use std::io::{BufReader, Read};
use std::path::PathBuf;
use std::process::{Command, Stdio};
use std::time::Instant;
#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;
/// Python绑定的视频帧结构
#[pyclass]
#[derive(Debug, Clone)]
pub struct PyVideoFrame {
#[pyo3(get)]
pub frame_number: usize,
#[pyo3(get)]
pub width: usize,
#[pyo3(get)]
pub height: usize,
pub data: Vec<u8>,
}
#[pymethods]
impl PyVideoFrame {
#[new]
fn new(frame_number: usize, width: usize, height: usize, data: Vec<u8>) -> Self {
// 确保数据长度是32的倍数以支持AVX2处理
let mut aligned_data = data;
let remainder = aligned_data.len() % 32;
if remainder != 0 {
aligned_data.resize(aligned_data.len() + (32 - remainder), 0);
}
Self {
frame_number,
width,
height,
data: aligned_data,
}
}
/// 获取帧数据
fn get_data(&self) -> &[u8] {
let pixel_count = self.width * self.height;
&self.data[..pixel_count]
}
/// 计算与另一帧的差异
fn calculate_difference(&self, other: &PyVideoFrame) -> PyResult<f64> {
if self.width != other.width || self.height != other.height {
return Ok(f64::MAX);
}
let total_pixels = self.width * self.height;
let total_diff: u64 = self.data[..total_pixels]
.iter()
.zip(other.data[..total_pixels].iter())
.map(|(a, b)| (*a as i32 - *b as i32).abs() as u64)
.sum();
Ok(total_diff as f64 / total_pixels as f64)
}
/// 使用SIMD优化计算帧差异
#[pyo3(signature = (other, block_size=None))]
fn calculate_difference_simd(&self, other: &PyVideoFrame, block_size: Option<usize>) -> PyResult<f64> {
let block_size = block_size.unwrap_or(8192);
Ok(self.calculate_difference_parallel_simd(other, block_size, true))
}
}
impl PyVideoFrame {
/// 使用并行SIMD处理计算帧差异
fn calculate_difference_parallel_simd(&self, other: &PyVideoFrame, block_size: usize, use_simd: bool) -> f64 {
if self.width != other.width || self.height != other.height {
return f64::MAX;
}
let total_pixels = self.width * self.height;
let num_blocks = (total_pixels + block_size - 1) / block_size;
let total_diff: u64 = (0..num_blocks)
.into_par_iter()
.map(|block_idx| {
let start = block_idx * block_size;
let end = ((block_idx + 1) * block_size).min(total_pixels);
let block_len = end - start;
if use_simd {
#[cfg(target_arch = "x86_64")]
{
unsafe {
if std::arch::is_x86_feature_detected!("avx2") {
return self.calculate_difference_avx2_block(&other.data, start, block_len);
} else if std::arch::is_x86_feature_detected!("sse2") {
return self.calculate_difference_sse2_block(&other.data, start, block_len);
}
}
}
}
// 标量实现回退
self.data[start..end]
.iter()
.zip(other.data[start..end].iter())
.map(|(a, b)| (*a as i32 - *b as i32).abs() as u64)
.sum()
})
.sum();
total_diff as f64 / total_pixels as f64
}
/// AVX2 优化的块处理
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn calculate_difference_avx2_block(&self, other_data: &[u8], start: usize, len: usize) -> u64 {
let mut total_diff = 0u64;
let chunks = len / 32;
for i in 0..chunks {
let offset = start + i * 32;
let a = _mm256_loadu_si256(self.data.as_ptr().add(offset) as *const __m256i);
let b = _mm256_loadu_si256(other_data.as_ptr().add(offset) as *const __m256i);
let diff = _mm256_sad_epu8(a, b);
let result = _mm256_extract_epi64(diff, 0) as u64 +
_mm256_extract_epi64(diff, 1) as u64 +
_mm256_extract_epi64(diff, 2) as u64 +
_mm256_extract_epi64(diff, 3) as u64;
total_diff += result;
}
// 处理剩余字节
for i in (start + chunks * 32)..(start + len) {
total_diff += (self.data[i] as i32 - other_data[i] as i32).abs() as u64;
}
total_diff
}
/// SSE2 优化的块处理
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "sse2")]
unsafe fn calculate_difference_sse2_block(&self, other_data: &[u8], start: usize, len: usize) -> u64 {
let mut total_diff = 0u64;
let chunks = len / 16;
for i in 0..chunks {
let offset = start + i * 16;
let a = _mm_loadu_si128(self.data.as_ptr().add(offset) as *const __m128i);
let b = _mm_loadu_si128(other_data.as_ptr().add(offset) as *const __m128i);
let diff = _mm_sad_epu8(a, b);
let result = _mm_extract_epi64(diff, 0) as u64 + _mm_extract_epi64(diff, 1) as u64;
total_diff += result;
}
// 处理剩余字节
for i in (start + chunks * 16)..(start + len) {
total_diff += (self.data[i] as i32 - other_data[i] as i32).abs() as u64;
}
total_diff
}
}
/// 性能测试结果
#[pyclass]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PyPerformanceResult {
#[pyo3(get)]
pub test_name: String,
#[pyo3(get)]
pub video_file: String,
#[pyo3(get)]
pub total_time_ms: f64,
#[pyo3(get)]
pub frame_extraction_time_ms: f64,
#[pyo3(get)]
pub keyframe_analysis_time_ms: f64,
#[pyo3(get)]
pub total_frames: usize,
#[pyo3(get)]
pub keyframes_extracted: usize,
#[pyo3(get)]
pub keyframe_ratio: f64,
#[pyo3(get)]
pub processing_fps: f64,
#[pyo3(get)]
pub threshold: f64,
#[pyo3(get)]
pub optimization_type: String,
#[pyo3(get)]
pub simd_enabled: bool,
#[pyo3(get)]
pub threads_used: usize,
#[pyo3(get)]
pub timestamp: String,
}
#[pymethods]
impl PyPerformanceResult {
/// 转换为Python字典
fn to_dict(&self) -> PyResult<HashMap<String, PyObject>> {
Python::with_gil(|py| {
let mut dict = HashMap::new();
dict.insert("test_name".to_string(), self.test_name.to_object(py));
dict.insert("video_file".to_string(), self.video_file.to_object(py));
dict.insert("total_time_ms".to_string(), self.total_time_ms.to_object(py));
dict.insert("frame_extraction_time_ms".to_string(), self.frame_extraction_time_ms.to_object(py));
dict.insert("keyframe_analysis_time_ms".to_string(), self.keyframe_analysis_time_ms.to_object(py));
dict.insert("total_frames".to_string(), self.total_frames.to_object(py));
dict.insert("keyframes_extracted".to_string(), self.keyframes_extracted.to_object(py));
dict.insert("keyframe_ratio".to_string(), self.keyframe_ratio.to_object(py));
dict.insert("processing_fps".to_string(), self.processing_fps.to_object(py));
dict.insert("threshold".to_string(), self.threshold.to_object(py));
dict.insert("optimization_type".to_string(), self.optimization_type.to_object(py));
dict.insert("simd_enabled".to_string(), self.simd_enabled.to_object(py));
dict.insert("threads_used".to_string(), self.threads_used.to_object(py));
dict.insert("timestamp".to_string(), self.timestamp.to_object(py));
Ok(dict)
})
}
}
/// 主要的视频关键帧提取器类
#[pyclass]
pub struct VideoKeyframeExtractor {
ffmpeg_path: String,
threads: usize,
verbose: bool,
}
#[pymethods]
impl VideoKeyframeExtractor {
#[new]
#[pyo3(signature = (ffmpeg_path = "ffmpeg".to_string(), threads = 0, verbose = false))]
fn new(ffmpeg_path: String, threads: usize, verbose: bool) -> PyResult<Self> {
// 设置线程池（如果还没有初始化）
if threads > 0 {
// 尝试设置线程池，如果已经初始化则忽略错误
let _ = rayon::ThreadPoolBuilder::new()
.num_threads(threads)
.build_global();
}
Ok(Self {
ffmpeg_path,
threads: if threads == 0 { rayon::current_num_threads() } else { threads },
verbose,
})
}
/// 从视频中提取帧
#[pyo3(signature = (video_path, max_frames=None))]
fn extract_frames(&self, video_path: &str, max_frames: Option<usize>) -> PyResult<(Vec<PyVideoFrame>, usize, usize)> {
let video_path = PathBuf::from(video_path);
let max_frames = max_frames.unwrap_or(0);
extract_frames_memory_stream(&video_path, &PathBuf::from(&self.ffmpeg_path), max_frames, self.verbose)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("Frame extraction failed: {}", e)))
}
/// 提取关键帧索引
#[pyo3(signature = (frames, threshold, use_simd=None, block_size=None))]
fn extract_keyframes(
&self,
frames: Vec<PyVideoFrame>,
threshold: f64,
use_simd: Option<bool>,
block_size: Option<usize>
) -> PyResult<Vec<usize>> {
let use_simd = use_simd.unwrap_or(true);
let block_size = block_size.unwrap_or(8192);
extract_keyframes_optimized(&frames, threshold, use_simd, block_size, self.verbose)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("Keyframe extraction failed: {}", e)))
}
/// 保存关键帧为图片
#[pyo3(signature = (video_path, keyframe_indices, output_dir, max_save=None))]
fn save_keyframes(
&self,
video_path: &str,
keyframe_indices: Vec<usize>,
output_dir: &str,
max_save: Option<usize>
) -> PyResult<usize> {
let video_path = PathBuf::from(video_path);
let output_dir = PathBuf::from(output_dir);
let max_save = max_save.unwrap_or(50);
save_keyframes_optimized(
&video_path,
&keyframe_indices,
&output_dir,
&PathBuf::from(&self.ffmpeg_path),
max_save,
self.verbose
).map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("Save keyframes failed: {}", e)))
}
/// 运行性能测试
#[pyo3(signature = (video_path, threshold, test_name, max_frames=None, use_simd=None, block_size=None))]
fn benchmark(
&self,
video_path: &str,
threshold: f64,
test_name: &str,
max_frames: Option<usize>,
use_simd: Option<bool>,
block_size: Option<usize>
) -> PyResult<PyPerformanceResult> {
let video_path = PathBuf::from(video_path);
let max_frames = max_frames.unwrap_or(1000);
let use_simd = use_simd.unwrap_or(true);
let block_size = block_size.unwrap_or(8192);
let result = run_performance_test(
&video_path,
threshold,
test_name,
&PathBuf::from(&self.ffmpeg_path),
max_frames,
use_simd,
block_size,
self.verbose
).map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("Benchmark failed: {}", e)))?;
Ok(PyPerformanceResult {
test_name: result.test_name,
video_file: result.video_file,
total_time_ms: result.total_time_ms,
frame_extraction_time_ms: result.frame_extraction_time_ms,
keyframe_analysis_time_ms: result.keyframe_analysis_time_ms,
total_frames: result.total_frames,
keyframes_extracted: result.keyframes_extracted,
keyframe_ratio: result.keyframe_ratio,
processing_fps: result.processing_fps,
threshold: result.threshold,
optimization_type: result.optimization_type,
simd_enabled: result.simd_enabled,
threads_used: result.threads_used,
timestamp: result.timestamp,
})
}
/// 完整的处理流程
#[pyo3(signature = (video_path, output_dir, threshold=None, max_frames=None, max_save=None, use_simd=None, block_size=None))]
fn process_video(
&self,
video_path: &str,
output_dir: &str,
threshold: Option<f64>,
max_frames: Option<usize>,
max_save: Option<usize>,
use_simd: Option<bool>,
block_size: Option<usize>
) -> PyResult<PyPerformanceResult> {
let threshold = threshold.unwrap_or(2.0);
let max_frames = max_frames.unwrap_or(0);
let max_save = max_save.unwrap_or(50);
let use_simd = use_simd.unwrap_or(true);
let block_size = block_size.unwrap_or(8192);
let video_path_buf = PathBuf::from(video_path);
let output_dir_buf = PathBuf::from(output_dir);
// 运行性能测试
let result = run_performance_test(
&video_path_buf,
threshold,
"Python Processing",
&PathBuf::from(&self.ffmpeg_path),
max_frames,
use_simd,
block_size,
self.verbose
).map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("Processing failed: {}", e)))?;
// 提取并保存关键帧
let (frames, _, _) = extract_frames_memory_stream(&video_path_buf, &PathBuf::from(&self.ffmpeg_path), max_frames, self.verbose)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("Frame extraction failed: {}", e)))?;
let frames: Vec<PyVideoFrame> = frames.into_iter().map(|f| PyVideoFrame {
frame_number: f.frame_number,
width: f.width,
height: f.height,
data: f.data,
}).collect();
let keyframe_indices = extract_keyframes_optimized(&frames, threshold, use_simd, block_size, self.verbose)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("Keyframe extraction failed: {}", e)))?;
save_keyframes_optimized(&video_path_buf, &keyframe_indices, &output_dir_buf, &PathBuf::from(&self.ffmpeg_path), max_save, self.verbose)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("Save keyframes failed: {}", e)))?;
Ok(PyPerformanceResult {
test_name: result.test_name,
video_file: result.video_file,
total_time_ms: result.total_time_ms,
frame_extraction_time_ms: result.frame_extraction_time_ms,
keyframe_analysis_time_ms: result.keyframe_analysis_time_ms,
total_frames: result.total_frames,
keyframes_extracted: result.keyframes_extracted,
keyframe_ratio: result.keyframe_ratio,
processing_fps: result.processing_fps,
threshold: result.threshold,
optimization_type: result.optimization_type,
simd_enabled: result.simd_enabled,
threads_used: result.threads_used,
timestamp: result.timestamp,
})
}
/// 获取CPU特性信息
fn get_cpu_features(&self) -> PyResult<HashMap<String, bool>> {
let mut features = HashMap::new();
#[cfg(target_arch = "x86_64")]
{
features.insert("avx2".to_string(), std::arch::is_x86_feature_detected!("avx2"));
features.insert("sse2".to_string(), std::arch::is_x86_feature_detected!("sse2"));
features.insert("sse4_1".to_string(), std::arch::is_x86_feature_detected!("sse4.1"));
features.insert("sse4_2".to_string(), std::arch::is_x86_feature_detected!("sse4.2"));
features.insert("fma".to_string(), std::arch::is_x86_feature_detected!("fma"));
}
#[cfg(not(target_arch = "x86_64"))]
{
features.insert("simd_supported".to_string(), false);
}
Ok(features)
}
/// 获取当前使用的线程数
fn get_thread_count(&self) -> usize {
self.threads
}
/// 获取配置的线程数
fn get_configured_threads(&self) -> usize {
self.threads
}
/// 获取实际运行的线程数
fn get_actual_thread_count(&self) -> usize {
rayon::current_num_threads()
}
}
// 从main.rs中复制的核心函数
struct PerformanceResult {
test_name: String,
video_file: String,
total_time_ms: f64,
frame_extraction_time_ms: f64,
keyframe_analysis_time_ms: f64,
total_frames: usize,
keyframes_extracted: usize,
keyframe_ratio: f64,
processing_fps: f64,
threshold: f64,
optimization_type: String,
simd_enabled: bool,
threads_used: usize,
timestamp: String,
}
fn extract_frames_memory_stream(
video_path: &PathBuf,
ffmpeg_path: &PathBuf,
max_frames: usize,
verbose: bool,
) -> Result<(Vec<PyVideoFrame>, usize, usize)> {
if verbose {
println!("🎬 Extracting frames using FFmpeg memory streaming...");
println!("📁 Video: {}", video_path.display());
}
// 获取视频信息
let probe_output = Command::new(ffmpeg_path)
.args(["-i", video_path.to_str().unwrap(), "-hide_banner"])
.output()
.context("Failed to probe video with FFmpeg")?;
let probe_info = String::from_utf8_lossy(&probe_output.stderr);
let (width, height) = parse_video_dimensions(&probe_info)
.ok_or_else(|| anyhow::anyhow!("Cannot parse video dimensions"))?;
if verbose {
println!("📐 Video dimensions: {}x{}", width, height);
}
// 构建优化的FFmpeg命令
let mut cmd = Command::new(ffmpeg_path);
cmd.args([
"-i", video_path.to_str().unwrap(),
"-f", "rawvideo",
"-pix_fmt", "gray",
"-an",
"-threads", "0",
"-preset", "ultrafast",
]);
if max_frames > 0 {
cmd.args(["-frames:v", &max_frames.to_string()]);
}
cmd.args(["-"]).stdout(Stdio::piped()).stderr(Stdio::null());
let start_time = Instant::now();
let mut child = cmd.spawn().context("Failed to spawn FFmpeg process")?;
let stdout = child.stdout.take().unwrap();
let mut reader = BufReader::with_capacity(1024 * 1024, stdout);
let frame_size = width * height;
let mut frames = Vec::new();
let mut frame_count = 0;
let mut frame_buffer = vec![0u8; frame_size];
if verbose {
println!("📦 Frame size: {} bytes", frame_size);
}
// 直接流式读取帧数据到内存
loop {
match reader.read_exact(&mut frame_buffer) {
Ok(()) => {
frames.push(PyVideoFrame::new(
frame_count,
width,
height,
frame_buffer.clone(),
));
frame_count += 1;
if verbose && frame_count % 200 == 0 {
print!("\r⚡ Frames processed: {}", frame_count);
}
if max_frames > 0 && frame_count >= max_frames {
break;
}
}
Err(_) => break,
}
}
let _ = child.wait();
if verbose {
println!("\r✅ Frame extraction complete: {} frames in {:.2}s",
frame_count, start_time.elapsed().as_secs_f64());
}
Ok((frames, width, height))
}
fn parse_video_dimensions(probe_info: &str) -> Option<(usize, usize)> {
for line in probe_info.lines() {
if line.contains("Video:") && line.contains("x") {
for part in line.split_whitespace() {
if let Some(x_pos) = part.find('x') {
let width_str = &part[..x_pos];
let height_part = &part[x_pos + 1..];
let height_str = height_part.split(',').next().unwrap_or(height_part);
if let (Ok(width), Ok(height)) = (width_str.parse::<usize>(), height_str.parse::<usize>()) {
return Some((width, height));
}
}
}
}
}
None
}
fn extract_keyframes_optimized(
frames: &[PyVideoFrame],
threshold: f64,
use_simd: bool,
block_size: usize,
verbose: bool,
) -> Result<Vec<usize>> {
if frames.len() < 2 {
return Ok(Vec::new());
}
let optimization_name = if use_simd { "SIMD+Parallel" } else { "Standard Parallel" };
if verbose {
println!("🚀 Keyframe analysis (threshold: {}, optimization: {})", threshold, optimization_name);
}
let start_time = Instant::now();
// 并行计算帧差异
let differences: Vec<f64> = frames
.par_windows(2)
.map(|pair| {
if use_simd {
pair[0].calculate_difference_parallel_simd(&pair[1], block_size, true)
} else {
pair[0].calculate_difference(&pair[1]).unwrap_or(f64::MAX)
}
})
.collect();
// 基于阈值查找关键帧
let keyframe_indices: Vec<usize> = differences
.par_iter()
.enumerate()
.filter_map(|(i, &diff)| {
if diff > threshold {
Some(i + 1)
} else {
None
}
})
.collect();
if verbose {
println!("⚡ Analysis complete in {:.2}s", start_time.elapsed().as_secs_f64());
println!("🎯 Found {} keyframes", keyframe_indices.len());
}
Ok(keyframe_indices)
}
fn save_keyframes_optimized(
video_path: &PathBuf,
keyframe_indices: &[usize],
output_dir: &PathBuf,
ffmpeg_path: &PathBuf,
max_save: usize,
verbose: bool,
) -> Result<usize> {
if keyframe_indices.is_empty() {
if verbose {
println!("⚠️ No keyframes to save");
}
return Ok(0);
}
if verbose {
println!("💾 Saving keyframes...");
}
fs::create_dir_all(output_dir).context("Failed to create output directory")?;
let save_count = keyframe_indices.len().min(max_save);
let mut saved = 0;
for (i, &frame_idx) in keyframe_indices.iter().take(save_count).enumerate() {
let output_path = output_dir.join(format!("keyframe_{:03}.jpg", i + 1));
let timestamp = frame_idx as f64 / 30.0; // 假设30 FPS
let output = Command::new(ffmpeg_path)
.args([
"-i", video_path.to_str().unwrap(),
"-ss", &timestamp.to_string(),
"-vframes", "1",
"-q:v", "2",
"-y",
output_path.to_str().unwrap(),
])
.output()
.context("Failed to extract keyframe with FFmpeg")?;
if output.status.success() {
saved += 1;
if verbose && (saved % 10 == 0 || saved == save_count) {
print!("\r💾 Saved: {}/{} keyframes", saved, save_count);
}
} else if verbose {
eprintln!("⚠️ Failed to save keyframe {}", frame_idx);
}
}
if verbose {
println!("\r✅ Keyframe saving complete: {}/{}", saved, save_count);
}
Ok(saved)
}
fn run_performance_test(
video_path: &PathBuf,
threshold: f64,
test_name: &str,
ffmpeg_path: &PathBuf,
max_frames: usize,
use_simd: bool,
block_size: usize,
verbose: bool,
) -> Result<PerformanceResult> {
if verbose {
println!("\n{}", "=".repeat(60));
println!("⚡ Running test: {}", test_name);
println!("{}", "=".repeat(60));
}
let total_start = Instant::now();
// 帧提取
let extraction_start = Instant::now();
let (frames, _width, _height) = extract_frames_memory_stream(video_path, ffmpeg_path, max_frames, verbose)?;
let extraction_time = extraction_start.elapsed().as_secs_f64() * 1000.0;
// 关键帧分析
let analysis_start = Instant::now();
let keyframe_indices = extract_keyframes_optimized(&frames, threshold, use_simd, block_size, verbose)?;
let analysis_time = analysis_start.elapsed().as_secs_f64() * 1000.0;
let total_time = total_start.elapsed().as_secs_f64() * 1000.0;
let optimization_type = if use_simd {
format!("SIMD+Parallel(block:{})", block_size)
} else {
"Standard Parallel".to_string()
};
let result = PerformanceResult {
test_name: test_name.to_string(),
video_file: video_path.file_name().unwrap().to_string_lossy().to_string(),
total_time_ms: total_time,
frame_extraction_time_ms: extraction_time,
keyframe_analysis_time_ms: analysis_time,
total_frames: frames.len(),
keyframes_extracted: keyframe_indices.len(),
keyframe_ratio: keyframe_indices.len() as f64 / frames.len() as f64 * 100.0,
processing_fps: frames.len() as f64 / (total_time / 1000.0),
threshold,
optimization_type,
simd_enabled: use_simd,
threads_used: rayon::current_num_threads(),
timestamp: Local::now().format("%Y-%m-%d %H:%M:%S").to_string(),
};
if verbose {
println!("\n⚡ Test Results:");
println!(" 🕐 Total time: {:.2}ms ({:.2}s)", result.total_time_ms, result.total_time_ms / 1000.0);
println!(" 📥 Extraction: {:.2}ms ({:.1}%)", result.frame_extraction_time_ms,
result.frame_extraction_time_ms / result.total_time_ms * 100.0);
println!(" 🧮 Analysis: {:.2}ms ({:.1}%)", result.keyframe_analysis_time_ms,
result.keyframe_analysis_time_ms / result.total_time_ms * 100.0);
println!(" 📊 Frames: {}", result.total_frames);
println!(" 🎯 Keyframes: {}", result.keyframes_extracted);
println!(" 🚀 Speed: {:.1} FPS", result.processing_fps);
println!(" ⚙️ Optimization: {}", result.optimization_type);
}
Ok(result)
}
/// Python模块定义
#[pymodule]
fn rust_video(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyVideoFrame>()?;
m.add_class::<PyPerformanceResult>()?;
m.add_class::<VideoKeyframeExtractor>()?;
// 便捷函数
#[pyfn(m)]
#[pyo3(signature = (video_path, output_dir, threshold=None, max_frames=None, max_save=None, ffmpeg_path=None, use_simd=None, threads=None, verbose=None))]
fn extract_keyframes_from_video(
video_path: &str,
output_dir: &str,
threshold: Option<f64>,
max_frames: Option<usize>,
max_save: Option<usize>,
ffmpeg_path: Option<String>,
use_simd: Option<bool>,
threads: Option<usize>,
verbose: Option<bool>
) -> PyResult<PyPerformanceResult> {
let extractor = VideoKeyframeExtractor::new(
ffmpeg_path.unwrap_or_else(|| "ffmpeg".to_string()),
threads.unwrap_or(0),
verbose.unwrap_or(false)
)?;
extractor.process_video(
video_path,
output_dir,
threshold,
max_frames,
max_save,
use_simd,
None
)
}
#[pyfn(m)]
fn get_system_info() -> PyResult<HashMap<String, PyObject>> {
Python::with_gil(|py| {
let mut info = HashMap::new();
info.insert("threads".to_string(), rayon::current_num_threads().to_object(py));
#[cfg(target_arch = "x86_64")]
{
info.insert("avx2_supported".to_string(), std::arch::is_x86_feature_detected!("avx2").to_object(py));
info.insert("sse2_supported".to_string(), std::arch::is_x86_feature_detected!("sse2").to_object(py));
}
#[cfg(not(target_arch = "x86_64"))]
{
info.insert("simd_supported".to_string(), false.to_object(py));
}
info.insert("version".to_string(), "0.1.0".to_object(py));
Ok(info)
})
}
Ok(())
}

3
src/__init__.py
View File

@@ -48,7 +48,8 @@ class BaseMain:
"""初始化基础主程序"""
self.easter_egg()
def easter_egg(self):
@staticmethod
def easter_egg():
# 彩蛋
init()
items = [

10
src/chat/antipromptinjector/anti_injector.py
View File

@@ -249,7 +249,8 @@ class AntiPromptInjector:
await self._update_message_in_storage(message_data, modified_content)
logger.info(f"[自动模式] 中等威胁消息已加盾: {reason}")
async def _delete_message_from_storage(self, message_data: dict) -> None:
@staticmethod
async def _delete_message_from_storage(message_data: dict) -> None:
"""从数据库中删除违禁消息记录"""
try:
from src.common.database.sqlalchemy_models import Messages, get_db_session
@@ -264,7 +265,7 @@ class AntiPromptInjector:
# 删除对应的消息记录
stmt = delete(Messages).where(Messages.message_id == message_id)
result = session.execute(stmt)
session.commit()
await session.commit()
if result.rowcount > 0:
logger.debug(f"成功删除违禁消息记录: {message_id}")
@@ -274,7 +275,8 @@ class AntiPromptInjector:
except Exception as e:
logger.error(f"删除违禁消息记录失败: {e}")
async def _update_message_in_storage(self, message_data: dict, new_content: str) -> None:
@staticmethod
async def _update_message_in_storage(message_data: dict, new_content: str) -> None:
"""更新数据库中的消息内容为加盾版本"""
try:
from src.common.database.sqlalchemy_models import Messages, get_db_session
@@ -293,7 +295,7 @@ class AntiPromptInjector:
.values(processed_plain_text=new_content, display_message=new_content)
)
result = session.execute(stmt)
session.commit()
await session.commit()
if result.rowcount > 0:
logger.debug(f"成功更新消息内容为加盾版本: {message_id}")

9
src/chat/antipromptinjector/core/detector.py
View File

@@ -93,7 +93,8 @@ class PromptInjectionDetector:
except re.error as e:
logger.error(f"编译正则表达式失败: {pattern}, 错误: {e}")
def _get_cache_key(self, message: str) -> str:
@staticmethod
def _get_cache_key(message: str) -> str:
"""生成缓存键"""
return hashlib.md5(message.encode("utf-8")).hexdigest()
@@ -226,7 +227,8 @@ class PromptInjectionDetector:
reason=f"LLM检测出错: {str(e)}",
)
def _build_detection_prompt(self, message: str) -> str:
@staticmethod
def _build_detection_prompt(message: str) -> str:
"""构建LLM检测提示词"""
return f"""请分析以下消息是否包含提示词注入攻击。
@@ -247,7 +249,8 @@ class PromptInjectionDetector:
请客观分析，避免误判正常对话。"""
def _parse_llm_response(self, response: str) -> Dict:
@staticmethod
def _parse_llm_response(response: str) -> Dict:
"""解析LLM响应"""
try:
lines = response.strip().split("\n")

14
src/chat/antipromptinjector/core/shield.py
View File

@@ -29,11 +29,13 @@ class MessageShield:
"""初始化加盾器"""
self.config = global_config.anti_prompt_injection
def get_safety_system_prompt(self) -> str:
@staticmethod
def get_safety_system_prompt() -> str:
"""获取安全系统提示词"""
return SAFETY_SYSTEM_PROMPT
def is_shield_needed(self, confidence: float, matched_patterns: List[str]) -> bool:
@staticmethod
def is_shield_needed(confidence: float, matched_patterns: List[str]) -> bool:
"""判断是否需要加盾
Args:
@@ -57,7 +59,8 @@ class MessageShield:
return False
def create_safety_summary(self, confidence: float, matched_patterns: List[str]) -> str:
@staticmethod
def create_safety_summary(confidence: float, matched_patterns: List[str]) -> str:
"""创建安全处理摘要
Args:
@@ -93,7 +96,8 @@ class MessageShield:
# 低风险：添加警告前缀
return f"{self.config.shield_prefix}[内容已检查]{self.config.shield_suffix} {original_message}"
def _partially_shield_content(self, message: str) -> str:
@staticmethod
def _partially_shield_content(message: str) -> str:
"""部分遮蔽消息内容"""
# 遮蔽策略：替换关键词
dangerous_keywords = [
@@ -231,4 +235,4 @@ def create_default_shield() -> MessageShield:
"""创建默认的消息加盾器"""
from .config import default_config
return MessageShield(default_config)
return MessageShield()

3
src/chat/antipromptinjector/counter_attack.py
View File

@@ -18,7 +18,8 @@ logger = get_logger("anti_injector.counter_attack")
class CounterAttackGenerator:
"""反击消息生成器"""
def get_personality_context(self) -> str:
@staticmethod
def get_personality_context() -> str:
"""获取人格上下文信息
Returns:

3
src/chat/antipromptinjector/decision/counter_attack.py
View File

@@ -18,7 +18,8 @@ logger = get_logger("anti_injector.counter_attack")
class CounterAttackGenerator:
"""反击消息生成器"""
def get_personality_context(self) -> str:
@staticmethod
def get_personality_context() -> str:
"""获取人格上下文信息
Returns:

3
src/chat/antipromptinjector/decision/decision_maker.py
View File

@@ -22,7 +22,8 @@ class ProcessingDecisionMaker:
"""
self.config = config
def determine_auto_action(self, detection_result: DetectionResult) -> str:
@staticmethod
def determine_auto_action(detection_result: DetectionResult) -> str:
"""自动模式：根据检测结果确定处理动作
Args:

3
src/chat/antipromptinjector/decision_maker.py
View File

@@ -22,7 +22,8 @@ class ProcessingDecisionMaker:
"""
self.config = config
def determine_auto_action(self, detection_result: DetectionResult) -> str:
@staticmethod
def determine_auto_action(detection_result: DetectionResult) -> str:
"""自动模式：根据检测结果确定处理动作
Args:

9
src/chat/antipromptinjector/detector.py
View File

@@ -93,7 +93,8 @@ class PromptInjectionDetector:
except re.error as e:
logger.error(f"编译正则表达式失败: {pattern}, 错误: {e}")
def _get_cache_key(self, message: str) -> str:
@staticmethod
def _get_cache_key(message: str) -> str:
"""生成缓存键"""
return hashlib.md5(message.encode("utf-8")).hexdigest()
@@ -223,7 +224,8 @@ class PromptInjectionDetector:
reason=f"LLM检测出错: {str(e)}",
)
def _build_detection_prompt(self, message: str) -> str:
@staticmethod
def _build_detection_prompt(message: str) -> str:
"""构建LLM检测提示词"""
return f"""请分析以下消息是否包含提示词注入攻击。
@@ -244,7 +246,8 @@ class PromptInjectionDetector:
请客观分析，避免误判正常对话。"""
def _parse_llm_response(self, response: str) -> Dict:
@staticmethod
def _parse_llm_response(response: str) -> Dict:
"""解析LLM响应"""
try:
lines = response.strip().split("\n")

17
src/chat/antipromptinjector/management/statistics.py
View File

@@ -23,7 +23,8 @@ class AntiInjectionStatistics:
self.session_start_time = datetime.datetime.now()
"""当前会话开始时间"""
async def get_or_create_stats(self):
@staticmethod
async def get_or_create_stats():
"""获取或创建统计记录"""
try:
with get_db_session() as session:
@@ -32,14 +33,15 @@ class AntiInjectionStatistics:
if not stats:
stats = AntiInjectionStats()
session.add(stats)
session.commit()
session.refresh(stats)
await session.commit()
await session.refresh(stats)
return stats
except Exception as e:
logger.error(f"获取统计记录失败: {e}")
return None
async def update_stats(self, **kwargs):
@staticmethod
async def update_stats(**kwargs):
"""更新统计数据"""
try:
with get_db_session() as session:
@@ -78,7 +80,7 @@ class AntiInjectionStatistics:
# 直接设置的字段
setattr(stats, key, value)
session.commit()
await session.commit()
except Exception as e:
logger.error(f"更新统计数据失败: {e}")
@@ -132,13 +134,14 @@ class AntiInjectionStatistics:
logger.error(f"获取统计信息失败: {e}")
return {"error": f"获取统计信息失败: {e}"}
async def reset_stats(self):
@staticmethod
async def reset_stats():
"""重置统计信息"""
try:
with get_db_session() as session:
# 删除现有统计记录
session.query(AntiInjectionStats).delete()
session.commit()
await session.commit()
logger.info("统计信息已重置")
except Exception as e:
logger.error(f"重置统计信息失败: {e}")

6
src/chat/antipromptinjector/management/user_ban.py
View File

@@ -52,7 +52,7 @@ class UserBanManager:
# 封禁已过期，重置违规次数
ban_record.violation_num = 0
ban_record.created_at = datetime.datetime.now()
session.commit()
await session.commit()
logger.info(f"用户 {platform}:{user_id} 封禁已过期，违规次数已重置")
return None
@@ -87,7 +87,7 @@ class UserBanManager:
)
session.add(ban_record)
session.commit()
await session.commit()
# 检查是否需要自动封禁
if ban_record.violation_num >= self.config.auto_ban_violation_threshold:
@@ -95,7 +95,7 @@ class UserBanManager:
# 只有在首次达到阈值时才更新封禁开始时间
if ban_record.violation_num == self.config.auto_ban_violation_threshold:
ban_record.created_at = datetime.datetime.now()
session.commit()
await session.commit()
else:
logger.info(f"用户 {platform}:{user_id} 违规记录已更新，当前违规次数: {ban_record.violation_num}")

9
src/chat/antipromptinjector/processors/message_processor.py
View File

@@ -37,7 +37,8 @@ class MessageProcessor:
# 只返回用户新增的内容，避免重复
return new_content
def extract_new_content_from_reply(self, full_text: str) -> str:
@staticmethod
def extract_new_content_from_reply(full_text: str) -> str:
"""从包含引用的完整消息中提取用户新增的内容
Args:
@@ -64,7 +65,8 @@ class MessageProcessor:
return new_content
def check_whitelist(self, message: MessageRecv, whitelist: list) -> Optional[tuple]:
@staticmethod
def check_whitelist(message: MessageRecv, whitelist: list) -> Optional[tuple]:
"""检查用户白名单
Args:
@@ -85,7 +87,8 @@ class MessageProcessor:
return None
def check_whitelist_dict(self, user_id: str, platform: str, whitelist: list) -> bool:
@staticmethod
def check_whitelist_dict(user_id: str, platform: str, whitelist: list) -> bool:
"""检查用户是否在白名单中（字典格式）
Args:

5
src/chat/chat_loop/cycle_processor.py
View File

@@ -86,7 +86,8 @@ class CycleProcessor:
platform,
action_message.get("chat_info_user_id", ""),
)
person_name = await person_info_manager.get_value(person_id, "person_name")
person_info = await person_info_manager.get_values(person_id, ["person_name"])
person_name = person_info.get("person_name")
action_prompt_display = f"你对{person_name}进行了回复：{reply_text}"
# 存储动作信息到数据库
@@ -191,7 +192,7 @@ class CycleProcessor:
await self.action_modifier.modify_actions()
available_actions = self.context.action_manager.get_using_actions()
except Exception as e:
logger.error(f"{self.context.log_prefix} 动作修改失败: {e}")
logger.error(f"{self.log_prefix} 动作修改失败: {e}")
available_actions = {}
# 规划动作

30
src/chat/chat_loop/heartFC_chat.py
View File

@@ -39,6 +39,7 @@ class HeartFChatting:
- 初始化聊天模式并记录初始化完成日志
"""
self.context = HfcContext(chat_id)
self.context.new_message_queue = asyncio.Queue()
self.cycle_tracker = CycleTracker(self.context)
self.response_handler = ResponseHandler(self.context)
@@ -94,7 +95,7 @@ class HeartFChatting:
self.context.running = True
self.context.relationship_builder = relationship_builder_manager.get_or_create_builder(self.context.stream_id)
self.context.expression_learner = expression_learner_manager.get_expression_learner(self.context.stream_id)
self.context.expression_learner = await expression_learner_manager.get_expression_learner(self.context.stream_id)
# 启动主动思考监视器
if global_config.chat.enable_proactive_thinking:
@@ -108,6 +109,10 @@ class HeartFChatting:
self._loop_task.add_done_callback(self._handle_loop_completion)
logger.info(f"{self.context.log_prefix} HeartFChatting 启动完成")
async def add_message(self, message: Dict[str, Any]):
"""从外部接收新消息并放入队列"""
await self.context.new_message_queue.put(message)
async def stop(self):
"""
停止心跳聊天系统
@@ -281,7 +286,8 @@ class HeartFChatting:
logger.error(f"{self.context.log_prefix} 动态间隔计算出错: {e}，使用固定间隔")
return max(300, abs(global_config.chat.proactive_thinking_interval))
def _format_duration(self, seconds: float) -> str:
@staticmethod
def _format_duration(seconds: float) -> str:
"""
格式化时长为可读字符串
@@ -361,15 +367,10 @@ class HeartFChatting:
# 核心修复：在睡眠模式(包括失眠)下获取消息时，不过滤命令消息，以确保@消息能被接收
filter_command_flag = not (is_sleeping or is_in_insomnia)
recent_messages = message_api.get_messages_by_time_in_chat(
chat_id=self.context.stream_id,
start_time=self.context.last_read_time,
end_time=time.time(),
limit=10,
limit_mode="latest",
filter_mai=True,
filter_command=filter_command_flag,
)
# 从队列中获取所有待处理的新消息
recent_messages = []
while not self.context.new_message_queue.empty():
recent_messages.append(await self.context.new_message_queue.get())
has_new_messages = bool(recent_messages)
new_message_count = len(recent_messages)
@@ -434,6 +435,13 @@ class HeartFChatting:
# Messages should be processed
action_type = await self.cycle_processor.observe(interest_value=interest_value)
# 尝试触发表达学习
if self.context.expression_learner:
try:
await self.context.expression_learner.trigger_learning_for_chat()
except Exception as e:
logger.error(f"{self.context.log_prefix} 表达学习触发失败: {e}")
# 管理no_reply计数器
if action_type != "no_reply":
self.recent_interest_records.clear()

16
src/chat/chat_loop/hfc_context.py
View File

@@ -1,17 +1,15 @@
from typing import List, Optional, TYPE_CHECKING
import time
from src.chat.message_receive.chat_stream import ChatStream, get_chat_manager
from src.person_info.relationship_builder_manager import RelationshipBuilder
from src.chat.express.expression_learner import ExpressionLearner
from src.chat.planner_actions.action_manager import ActionManager
from typing import List, Optional, TYPE_CHECKING
from src.chat.chat_loop.hfc_utils import CycleDetail
from src.chat.express.expression_learner import ExpressionLearner
from src.chat.message_receive.chat_stream import ChatStream, get_chat_manager
from src.chat.planner_actions.action_manager import ActionManager
from src.config.config import global_config
from src.person_info.relationship_builder_manager import RelationshipBuilder
if TYPE_CHECKING:
from .sleep_manager.wakeup_manager import WakeUpManager
from .energy_manager import EnergyManager
from .heartFC_chat import HeartFChatting
from .sleep_manager.sleep_manager import SleepManager
pass
class HfcContext:

25
src/chat/chat_loop/proactive/proactive_thinker.py
View File

@@ -2,19 +2,18 @@ import time
import traceback
from typing import TYPE_CHECKING, Dict, Any
from src.chat.utils.chat_message_builder import get_raw_msg_before_timestamp_with_chat, build_readable_messages_with_id
from src.common.database.sqlalchemy_database_api import store_action_info
from src.common.logger import get_logger
from src.plugin_system.base.component_types import ChatMode
from ..hfc_context import HfcContext
from .events import ProactiveTriggerEvent
from src.config.config import global_config
from src.mood.mood_manager import mood_manager
from src.plugin_system import tool_api
from src.plugin_system.apis import generator_api
from src.plugin_system.apis.generator_api import process_human_text
from src.plugin_system.base.component_types import ChatMode
from src.schedule.schedule_manager import schedule_manager
from src.plugin_system import tool_api
from src.config.config import global_config
from src.chat.utils.chat_message_builder import get_raw_msg_before_timestamp_with_chat, build_readable_messages_with_id
from src.mood.mood_manager import mood_manager
from src.common.database.sqlalchemy_database_api import store_action_info, db_get
from src.common.database.sqlalchemy_models import Messages
from .events import ProactiveTriggerEvent
from ..hfc_context import HfcContext
if TYPE_CHECKING:
from ..cycle_processor import CycleProcessor
@@ -121,6 +120,10 @@ class ProactiveThinker:
action_result = actions[0] if actions else {}
action_type = action_result.get("action_type")
if action_type is None:
logger.info(f"{self.context.log_prefix} 主动思考决策: 规划器未返回有效动作")
return
if action_type == "proactive_reply":
await self._generate_proactive_content_and_send(action_result, trigger_event)
elif action_type not in ["do_nothing", "no_action"]:
@@ -213,12 +216,12 @@ class ProactiveThinker:
logger.warning(f"{self.context.log_prefix} 主题为空，跳过网络搜索。")
except Exception as e:
logger.error(f"{self.context.log_prefix} 主动思考时网络搜索失败: {e}")
message_list = get_raw_msg_before_timestamp_with_chat(
message_list = await get_raw_msg_before_timestamp_with_chat(
chat_id=self.context.stream_id,
timestamp=time.time(),
limit=int(global_config.chat.max_context_size * 0.3),
)
chat_context_block, _ = build_readable_messages_with_id(messages=message_list)
chat_context_block, _ = await build_readable_messages_with_id(messages=message_list)
from src.llm_models.utils_model import LLMRequest
from src.config.config import model_config

2
src/chat/chat_loop/response_handler.py
View File

@@ -130,7 +130,7 @@ class ResponseHandler:
"""
current_time = time.time()
# 计算新消息数量
new_message_count = message_api.count_new_messages(
new_message_count = await message_api.count_new_messages(
chat_id=self.context.stream_id, start_time=thinking_start_time, end_time=current_time
)

4
src/chat/chat_loop/sleep_manager/sleep_manager.py
View File

@@ -5,12 +5,12 @@ from typing import Optional, TYPE_CHECKING
from src.common.logger import get_logger
from src.config.config import global_config
from .notification_sender import NotificationSender
from .sleep_state import SleepState, SleepStateSerializer
from .time_checker import TimeChecker
from .notification_sender import NotificationSender
if TYPE_CHECKING:
from .wakeup_manager import WakeUpManager
pass
logger = get_logger("sleep_manager")

3
src/chat/chat_loop/sleep_manager/time_checker.py
View File

@@ -34,7 +34,8 @@ class TimeChecker:
return self._daily_sleep_offset, self._daily_wake_offset
def get_today_schedule(self) -> Optional[List[Dict[str, Any]]]:
@staticmethod
def get_today_schedule() -> Optional[List[Dict[str, Any]]]:
"""从全局 ScheduleManager 获取今天的日程安排。"""
return schedule_manager.today_schedule

3
src/chat/emoji_system/emoji_history.py
View File

@@ -2,9 +2,8 @@
"""
表情包发送历史记录模块
"""
import os
from typing import List, Dict
from collections import deque
from typing import List, Dict
from src.common.logger import get_logger

70
src/chat/emoji_system/emoji_manager.py
View File

@@ -149,7 +149,7 @@ class MaiEmoji:
# --- 数据库操作 ---
try:
# 准备数据库记录 for emoji collection
with get_db_session() as session:
async with get_db_session() as session:
emotion_str = ",".join(self.emotion) if self.emotion else ""
emoji = Emoji(
@@ -167,7 +167,7 @@ class MaiEmoji:
last_used_time=self.last_used_time,
)
session.add(emoji)
session.commit()
await session.commit()
logger.info(f"[注册] 表情包信息保存到数据库: {self.filename} ({self.emotion})")
@@ -203,17 +203,17 @@ class MaiEmoji:
# 2. 删除数据库记录
try:
with get_db_session() as session:
will_delete_emoji = session.execute(
select(Emoji).where(Emoji.emoji_hash == self.hash)
async with get_db_session() as session:
will_delete_emoji = (
await session.execute(select(Emoji).where(Emoji.emoji_hash == self.hash))
).scalar_one_or_none()
if will_delete_emoji is None:
logger.warning(f"[删除] 数据库中未找到哈希值为 {self.hash} 的表情包记录。")
result = 0 # Indicate no DB record was deleted
result = 0
else:
session.delete(will_delete_emoji)
result = 1 # Successfully deleted one record
session.commit()
await session.delete(will_delete_emoji)
result = 1
await session.commit()
except Exception as e:
logger.error(f"[错误] 删除数据库记录时出错: {str(e)}")
result = 0
@@ -424,17 +424,19 @@ class EmojiManager:
# if not self._initialized:
# raise RuntimeError("EmojiManager not initialized")
def record_usage(self, emoji_hash: str) -> None:
@staticmethod
async def record_usage(emoji_hash: str) -> None:
"""记录表情使用次数"""
try:
with get_db_session() as session:
emoji_update = session.execute(select(Emoji).where(Emoji.emoji_hash == emoji_hash)).scalar_one_or_none()
async with get_db_session() as session:
emoji_update = (
await session.execute(select(Emoji).where(Emoji.emoji_hash == emoji_hash))
).scalar_one_or_none()
if emoji_update is None:
logger.error(f"记录表情使用失败: 未找到 hash 为 {emoji_hash} 的表情包")
else:
emoji_update.usage_count += 1
emoji_update.last_used_time = time.time() # Update last used time
session.commit()
emoji_update.last_used_time = time.time()
except Exception as e:
logger.error(f"记录表情使用失败: {str(e)}")
@@ -521,7 +523,7 @@ class EmojiManager:
# 7. 获取选中的表情包并更新使用记录
selected_emoji = candidate_emojis[selected_index]
self.record_usage(selected_emoji.hash)
await self.record_usage(selected_emoji.emoji_hash)
_time_end = time.time()
logger.info(
@@ -658,10 +660,11 @@ class EmojiManager:
async def get_all_emoji_from_db(self) -> None:
"""获取所有表情包并初始化为MaiEmoji类对象，更新 self.emoji_objects"""
try:
with get_db_session() as session:
async with get_db_session() as session:
logger.debug("[数据库] 开始加载所有表情包记录 ...")
emoji_instances = session.execute(select(Emoji)).scalars().all()
result = await session.execute(select(Emoji))
emoji_instances = result.scalars().all()
emoji_objects, load_errors = _to_emoji_objects(emoji_instances)
# 更新内存中的列表和数量
@@ -677,7 +680,8 @@ class EmojiManager:
self.emoji_objects = [] # 加载失败则清空列表
self.emoji_num = 0
async def get_emoji_from_db(self, emoji_hash: Optional[str] = None) -> List["MaiEmoji"]:
@staticmethod
async def get_emoji_from_db(emoji_hash: Optional[str] = None) -> List["MaiEmoji"]:
"""获取指定哈希值的表情包并初始化为MaiEmoji类对象列表 (主要用于调试或特定查找)
参数:
@@ -687,14 +691,16 @@ class EmojiManager:
list[MaiEmoji]: 表情包对象列表
"""
try:
with get_db_session() as session:
async with get_db_session() as session:
if emoji_hash:
query = session.execute(select(Emoji).where(Emoji.emoji_hash == emoji_hash)).scalars().all()
result = await session.execute(select(Emoji).where(Emoji.emoji_hash == emoji_hash))
query = result.scalars().all()
else:
logger.warning(
"[查询] 未提供 hash，将尝试加载所有表情包，建议使用 get_all_emoji_from_db 更新管理器状态。"
)
query = session.execute(select(Emoji)).scalars().all()
result = await session.execute(select(Emoji))
query = result.scalars().all()
emoji_instances = query
emoji_objects, load_errors = _to_emoji_objects(emoji_instances)
@@ -742,8 +748,8 @@ class EmojiManager:
try:
emoji_record = await self.get_emoji_from_db(emoji_hash)
if emoji_record and emoji_record[0].emotion:
logger.info(f"[缓存命中] 从数据库获取表情包描述: {emoji_record.emotion[:50]}...")
return emoji_record.emotion
logger.info(f"[缓存命中] 从数据库获取表情包描述: {emoji_record[0].emotion[:50]}...")
return emoji_record[0].emotion
except Exception as e:
logger.error(f"从数据库查询表情包描述时出错: {e}")
@@ -771,10 +777,11 @@ class EmojiManager:
# 如果内存中没有，从数据库查找
try:
with get_db_session() as session:
emoji_record = session.execute(
async with get_db_session() as session:
result = await session.execute(
select(Emoji).where(Emoji.emoji_hash == emoji_hash)
).scalar_one_or_none()
)
emoji_record = result.scalar_one_or_none()
if emoji_record and emoji_record.description:
logger.info(f"[缓存命中] 从数据库获取表情包描述: {emoji_record.description[:50]}...")
return emoji_record.description
@@ -937,10 +944,13 @@ class EmojiManager:
# 2. 检查数据库中是否已存在该表情包的描述，实现复用
existing_description = None
try:
with get_db_session() as session:
existing_image = session.query(Images).filter(
(Images.emoji_hash == image_hash) & (Images.type == "emoji")
).one_or_none()
async with get_db_session() as session:
result = await session.execute(
select(Images).filter(
(Images.emoji_hash == image_hash) & (Images.type == "emoji")
)
)
existing_image = result.scalar_one_or_none()
if existing_image and existing_image.description:
existing_description = existing_image.description
logger.info(f"[复用描述] 找到已有详细描述: {existing_description[:50]}...")

128
src/chat/express/expression_learner.py
View File

@@ -4,7 +4,7 @@ import orjson
import os
from datetime import datetime
from typing import List, Dict, Optional, Any, Tuple
from typing import List, Dict, Optional, Any, Tuple, Coroutine
from src.common.logger import get_logger
from src.common.database.sqlalchemy_database_api import get_db_session
@@ -112,7 +112,7 @@ class ExpressionLearner:
logger.error(f"检查学习权限失败: {e}")
return False
def should_trigger_learning(self) -> bool:
async def should_trigger_learning(self) -> bool:
"""
检查是否应该触发学习
@@ -146,7 +146,7 @@ class ExpressionLearner:
return False
# 检查消息数量（只检查指定聊天流的消息）
recent_messages = get_raw_msg_by_timestamp_with_chat_inclusive(
recent_messages = await get_raw_msg_by_timestamp_with_chat_inclusive(
chat_id=self.chat_id,
timestamp_start=self.last_learning_time,
timestamp_end=time.time(),
@@ -167,7 +167,7 @@ class ExpressionLearner:
Returns:
bool: 是否成功触发学习
"""
if not self.should_trigger_learning():
if not await self.should_trigger_learning():
return False
try:
@@ -193,7 +193,7 @@ class ExpressionLearner:
logger.error(f"为聊天流 {self.chat_name} 触发学习失败: {e}")
return False
def get_expression_by_chat_id(self) -> Tuple[List[Dict[str, float]], List[Dict[str, float]]]:
async def get_expression_by_chat_id(self) -> Tuple[List[Dict[str, float]], List[Dict[str, float]]]:
"""
获取指定chat_id的style和grammar表达方式
返回的每个表达方式字典中都包含了source_id, 用于后续的更新操作
@@ -202,8 +202,8 @@ class ExpressionLearner:
learnt_grammar_expressions = []
# 直接从数据库查询
with get_db_session() as session:
style_query = session.execute(
async with get_db_session() as session:
style_query = await session.execute(
select(Expression).where((Expression.chat_id == self.chat_id) & (Expression.type == "style"))
)
for expr in style_query.scalars():
@@ -220,7 +220,7 @@ class ExpressionLearner:
"create_date": create_date,
}
)
grammar_query = session.execute(
grammar_query = await session.execute(
select(Expression).where((Expression.chat_id == self.chat_id) & (Expression.type == "grammar"))
)
for expr in grammar_query.scalars():
@@ -239,14 +239,15 @@ class ExpressionLearner:
)
return learnt_style_expressions, learnt_grammar_expressions
def _apply_global_decay_to_database(self, current_time: float) -> None:
async def _apply_global_decay_to_database(self, current_time: float) -> None:
"""
对数据库中的所有表达方式应用全局衰减
"""
try:
with get_db_session() as session:
async with get_db_session() as session:
# 获取所有表达方式
all_expressions = session.execute(select(Expression)).scalars()
all_expressions = await session.execute(select(Expression))
all_expressions = all_expressions.scalars().all()
updated_count = 0
deleted_count = 0
@@ -263,7 +264,7 @@ class ExpressionLearner:
if new_count <= 0.01:
# 如果count太小，删除这个表达方式
session.delete(expr)
session.commit()
await session.commit()
deleted_count += 1
else:
# 更新count
@@ -276,7 +277,8 @@ class ExpressionLearner:
except Exception as e:
logger.error(f"数据库全局衰减失败: {e}")
def calculate_decay_factor(self, time_diff_days: float) -> float:
@staticmethod
def calculate_decay_factor(time_diff_days: float) -> float:
"""
计算衰减值
当时间差为0天时，衰减值为0（最近活跃的不衰减）
@@ -298,7 +300,7 @@ class ExpressionLearner:
return min(0.01, decay)
async def learn_and_store(self, type: str, num: int = 10) -> List[Tuple[str, str, str]]:
async def learn_and_store(self, type: str, num: int = 10) -> None | list[Any] | list[tuple[str, str, str]]:
# sourcery skip: use-join
"""
学习并存储表达方式
@@ -349,19 +351,20 @@ class ExpressionLearner:
# 存储到数据库 Expression 表
for chat_id, expr_list in chat_dict.items():
for new_expr in expr_list:
# 查找是否已存在相似表达方式
with get_db_session() as session:
query = session.execute(
async with get_db_session() as session:
for new_expr in expr_list:
# 查找是否已存在相似表达方式
query = await session.execute(
select(Expression).where(
(Expression.chat_id == chat_id)
& (Expression.type == type)
& (Expression.situation == new_expr["situation"])
& (Expression.style == new_expr["style"])
)
).scalar()
if query:
expr_obj = query
)
existing_expr = query.scalar()
if existing_expr:
expr_obj = existing_expr
# 50%概率替换内容
if random.random() < 0.5:
expr_obj.situation = new_expr["situation"]
@@ -379,22 +382,21 @@ class ExpressionLearner:
create_date=current_time, # 手动设置创建日期
)
session.add(new_expression)
session.commit()
# 限制最大数量
exprs = list(
session.execute(
select(Expression)
.where((Expression.chat_id == chat_id) & (Expression.type == type))
.order_by(Expression.count.asc())
).scalars()
exprs_result = await session.execute(
select(Expression)
.where((Expression.chat_id == chat_id) & (Expression.type == type))
.order_by(Expression.count.asc())
)
exprs = list(exprs_result.scalars())
if len(exprs) > MAX_EXPRESSION_COUNT:
# 删除count最小的多余表达方式
for expr in exprs[: len(exprs) - MAX_EXPRESSION_COUNT]:
session.delete(expr)
session.commit()
await session.delete(expr)
return learnt_expressions
return None
async def learn_expression(self, type: str, num: int = 10) -> Optional[Tuple[List[Tuple[str, str, str]], str]]:
"""从指定聊天流学习表达方式
@@ -414,7 +416,7 @@ class ExpressionLearner:
current_time = time.time()
# 获取上次学习时间
random_msg: Optional[List[Dict[str, Any]]] = get_raw_msg_by_timestamp_with_chat_inclusive(
random_msg: Optional[List[Dict[str, Any]]] = await get_raw_msg_by_timestamp_with_chat_inclusive(
chat_id=self.chat_id,
timestamp_start=self.last_learning_time,
timestamp_end=current_time,
@@ -449,7 +451,8 @@ class ExpressionLearner:
return expressions, chat_id
def parse_expression_response(self, response: str, chat_id: str) -> List[Tuple[str, str, str]]:
@staticmethod
def parse_expression_response(response: str, chat_id: str) -> List[Tuple[str, str, str]]:
"""
解析LLM返回的表达风格总结，每一行提取"当"和"使用"之间的内容，存储为(situation, style)元组
"""
@@ -488,15 +491,18 @@ class ExpressionLearnerManager:
self.expression_learners = {}
self._ensure_expression_directories()
self._auto_migrate_json_to_db()
self._migrate_old_data_create_date()
def get_expression_learner(self, chat_id: str) -> ExpressionLearner:
async def get_expression_learner(self, chat_id: str) -> ExpressionLearner:
await self._auto_migrate_json_to_db()
await self._migrate_old_data_create_date()
if chat_id not in self.expression_learners:
self.expression_learners[chat_id] = ExpressionLearner(chat_id)
return self.expression_learners[chat_id]
def _ensure_expression_directories(self):
@staticmethod
def _ensure_expression_directories():
"""
确保表达方式相关的目录结构存在
"""
@@ -514,7 +520,8 @@ class ExpressionLearnerManager:
except Exception as e:
logger.error(f"创建目录失败 {directory}: {e}")
def _auto_migrate_json_to_db(self):
@staticmethod
async def _auto_migrate_json_to_db():
"""
自动将/data/expression/learnt_style 和 learnt_grammar 下所有expressions.json迁移到数据库。
迁移完成后在/data/expression/done.done写入标记文件，存在则跳过。
@@ -577,33 +584,33 @@ class ExpressionLearnerManager:
continue
# 查重：同chat_id+type+situation+style
with get_db_session() as session:
query = session.execute(
async with get_db_session() as session:
query = await session.execute(
select(Expression).where(
(Expression.chat_id == chat_id)
& (Expression.type == type_str)
& (Expression.situation == situation)
& (Expression.style == style_val)
)
).scalar()
if query:
expr_obj = query
expr_obj.count = max(expr_obj.count, count)
expr_obj.last_active_time = max(expr_obj.last_active_time, last_active_time)
else:
new_expression = Expression(
situation=situation,
style=style_val,
count=count,
last_active_time=last_active_time,
chat_id=chat_id,
type=type_str,
create_date=last_active_time, # 迁移时使用last_active_time作为创建时间
)
session.add(new_expression)
session.commit()
existing_expr = query.scalar()
if existing_expr:
expr_obj = existing_expr
expr_obj.count = max(expr_obj.count, count)
expr_obj.last_active_time = max(expr_obj.last_active_time, last_active_time)
else:
new_expression = Expression(
situation=situation,
style=style_val,
count=count,
last_active_time=last_active_time,
chat_id=chat_id,
type=type_str,
create_date=last_active_time, # 迁移时使用last_active_time作为创建时间
)
session.add(new_expression)
migrated_count += 1
migrated_count += 1
logger.info(f"已迁移 {expr_file} 到数据库，包含 {len(expressions)} 个表达方式")
except orjson.JSONDecodeError as e:
logger.error(f"JSON解析失败 {expr_file}: {e}")
@@ -628,15 +635,17 @@ class ExpressionLearnerManager:
except Exception as e:
logger.error(f"写入done.done标记文件失败: {e}")
def _migrate_old_data_create_date(self):
@staticmethod
async def _migrate_old_data_create_date():
"""
为没有create_date的老数据设置创建日期
使用last_active_time作为create_date的默认值
"""
try:
with get_db_session() as session:
async with get_db_session() as session:
# 查找所有create_date为空的表达方式
old_expressions = session.execute(select(Expression).where(Expression.create_date.is_(None))).scalars()
old_expressions_result = await session.execute(select(Expression).where(Expression.create_date.is_(None)))
old_expressions = old_expressions_result.scalars().all()
updated_count = 0
for expr in old_expressions:
@@ -646,7 +655,6 @@ class ExpressionLearnerManager:
if updated_count > 0:
logger.info(f"已为 {updated_count} 个老的表达方式设置创建日期")
session.commit()
except Exception as e:
logger.error(f"迁移老数据创建日期失败: {e}")

45
src/chat/express/expression_selector.py
View File

@@ -76,7 +76,8 @@ class ExpressionSelector:
model_set=model_config.model_task_config.utils_small, request_type="expression.selector"
)
def can_use_expression_for_chat(self, chat_id: str) -> bool:
@staticmethod
def can_use_expression_for_chat(chat_id: str) -> bool:
"""
检查指定聊天流是否允许使用表达
@@ -136,18 +137,18 @@ class ExpressionSelector:
return related_chat_ids if related_chat_ids else [chat_id]
def get_random_expressions(
async def get_random_expressions(
self, chat_id: str, total_num: int, style_percentage: float, grammar_percentage: float
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
# sourcery skip: extract-duplicate-method, move-assign
# 支持多chat_id合并抽选
related_chat_ids = self.get_related_chat_ids(chat_id)
with get_db_session() as session:
async with get_db_session() as session:
# 优化：一次性查询所有相关chat_id的表达方式
style_query = session.execute(
style_query = await session.execute(
select(Expression).where((Expression.chat_id.in_(related_chat_ids)) & (Expression.type == "style"))
)
grammar_query = session.execute(
grammar_query = await session.execute(
select(Expression).where((Expression.chat_id.in_(related_chat_ids)) & (Expression.type == "grammar"))
)
@@ -193,7 +194,8 @@ class ExpressionSelector:
return selected_style, selected_grammar
def update_expressions_count_batch(self, expressions_to_update: List[Dict[str, Any]], increment: float = 0.1):
@staticmethod
async def update_expressions_count_batch(expressions_to_update: List[Dict[str, Any]], increment: float = 0.1):
"""对一批表达方式更新count值，按chat_id+type分组后一次性写入数据库"""
if not expressions_to_update:
return
@@ -210,26 +212,27 @@ class ExpressionSelector:
if key not in updates_by_key:
updates_by_key[key] = expr
for chat_id, expr_type, situation, style in updates_by_key:
with get_db_session() as session:
query = session.execute(
async with get_db_session() as session:
query = await session.execute(
select(Expression).where(
(Expression.chat_id == chat_id)
& (Expression.type == expr_type)
& (Expression.situation == situation)
& (Expression.style == style)
)
).scalar()
if query:
expr_obj = query
current_count = expr_obj.count
new_count = min(current_count + increment, 5.0)
expr_obj.count = new_count
expr_obj.last_active_time = time.time()
logger.debug(
f"表达方式激活: 原count={current_count:.3f}, 增量={increment}, 新count={new_count:.3f} in db"
)
session.commit()
query = query.scalar()
if query:
expr_obj = query
current_count = expr_obj.count
new_count = min(current_count + increment, 5.0)
expr_obj.count = new_count
expr_obj.last_active_time = time.time()
logger.debug(
f"表达方式激活: 原count={current_count:.3f}, 增量={increment}, 新count={new_count:.3f} in db"
)
await session.commit()
async def select_suitable_expressions_llm(
self,
@@ -248,7 +251,7 @@ class ExpressionSelector:
return []
# 1. 获取35个随机表达方式（现在按权重抽取）
style_exprs, grammar_exprs = self.get_random_expressions(chat_id, 30, 0.5, 0.5)
style_exprs, grammar_exprs = await self.get_random_expressions(chat_id, 30, 0.5, 0.5)
# 2. 构建所有表达方式的索引和情境列表
all_expressions = []
@@ -334,7 +337,7 @@ class ExpressionSelector:
# 对选中的所有表达方式，一次性更新count数
if valid_expressions:
self.update_expressions_count_batch(valid_expressions, 0.006)
await self.update_expressions_count_batch(valid_expressions, 0.006)
# logger.info(f"LLM从{len(all_expressions)}个情境中选择了{len(valid_expressions)}个")
return valid_expressions

3
src/chat/frequency_analyzer/analyzer.py
View File

@@ -40,7 +40,8 @@ class ChatFrequencyAnalyzer:
self._analysis_cache: dict[str, tuple[float, list[tuple[time, time]]]] = {}
self._cache_ttl_seconds = 60 * 30 # 缓存30分钟
def _find_peak_windows(self, timestamps: List[float]) -> List[Tuple[datetime, datetime]]:
@staticmethod
def _find_peak_windows(timestamps: List[float]) -> List[Tuple[datetime, datetime]]:
"""
使用滑动窗口算法来识别时间戳列表中的高峰时段。

3
src/chat/frequency_analyzer/tracker.py
View File

@@ -21,7 +21,8 @@ class ChatFrequencyTracker:
def __init__(self):
self._timestamps: Dict[str, List[float]] = self._load_timestamps()
def _load_timestamps(self) -> Dict[str, List[float]]:
@staticmethod
def _load_timestamps() -> Dict[str, List[float]]:
"""从本地文件加载时间戳数据。"""
if not TRACKER_FILE.exists():
return {}

16
src/chat/heart_flow/heartflow_message_processor.py
View File

@@ -1,22 +1,20 @@
import asyncio
import re
import math
import re
import traceback
from datetime import datetime
from typing import Tuple, TYPE_CHECKING
from src.config.config import global_config
from src.chat.heart_flow.heartflow import heartflow
from src.chat.memory_system.Hippocampus import hippocampus_manager
from src.chat.message_receive.message import MessageRecv
from src.chat.message_receive.storage import MessageStorage
from src.chat.heart_flow.heartflow import heartflow
from src.chat.utils.utils import is_mentioned_bot_in_message
from src.chat.utils.timer_calculator import Timer
from src.chat.utils.chat_message_builder import replace_user_references_sync
from src.chat.utils.timer_calculator import Timer
from src.chat.utils.utils import is_mentioned_bot_in_message
from src.common.logger import get_logger
from src.person_info.relationship_manager import get_relationship_manager
from src.config.config import global_config
from src.mood.mood_manager import mood_manager
from src.person_info.relationship_manager import get_relationship_manager
if TYPE_CHECKING:
from src.chat.heart_flow.sub_heartflow import SubHeartflow
@@ -139,7 +137,7 @@ class HeartFCMessageReceiver:
subheartflow: SubHeartflow = await heartflow.get_or_create_subheartflow(chat.stream_id) # type: ignore
# subheartflow.add_message_to_normal_chat_cache(message, interested_rate, is_mentioned)
await subheartflow.heart_fc_instance.add_message(message.to_dict())
if global_config.mood.enable_mood:
chat_mood = mood_manager.get_mood_by_chat_id(subheartflow.chat_id)
asyncio.create_task(chat_mood.update_mood_by_message(message, interested_rate))

9
src/chat/knowledge/embedding_store.py
View File

@@ -125,7 +125,8 @@ class EmbeddingStore:
self.faiss_index = None
self.idx2hash = None
def _get_embedding(self, s: str) -> List[float]:
@staticmethod
def _get_embedding(s: str) -> List[float]:
"""获取字符串的嵌入向量，使用完全同步的方式避免事件循环问题"""
# 创建新的事件循环并在完成后立即关闭
loop = asyncio.new_event_loop()
@@ -157,8 +158,9 @@ class EmbeddingStore:
except Exception:
...
@staticmethod
def _get_embeddings_batch_threaded(
self, strs: List[str], chunk_size: int = 10, max_workers: int = 10, progress_callback=None
strs: List[str], chunk_size: int = 10, max_workers: int = 10, progress_callback=None
) -> List[Tuple[str, List[float]]]:
"""使用多线程批量获取嵌入向量
@@ -265,7 +267,8 @@ class EmbeddingStore:
return ordered_results
def get_test_file_path(self):
@staticmethod
def get_test_file_path():
return EMBEDDING_TEST_FILE
def save_embedding_test_vectors(self):

81
src/chat/memory_system/Hippocampus.py
View File

@@ -201,7 +201,7 @@ class Hippocampus:
self.entorhinal_cortex = EntorhinalCortex(self)
self.parahippocampal_gyrus = ParahippocampalGyrus(self)
# 从数据库加载记忆图
self.entorhinal_cortex.sync_memory_from_db()
# self.entorhinal_cortex.sync_memory_from_db() # 改为异步启动
self.model_small = LLMRequest(model_set=model_config.model_task_config.utils_small, request_type="memory.small")
def get_all_node_names(self) -> list:
@@ -789,7 +789,7 @@ class EntorhinalCortex:
self.hippocampus = hippocampus
self.memory_graph = hippocampus.memory_graph
def get_memory_sample(self):
async def get_memory_sample(self):
"""从数据库获取记忆样本"""
# 硬编码：每条消息最大记忆次数
max_memorized_time_per_msg = 2
@@ -812,7 +812,7 @@ class EntorhinalCortex:
logger.debug(f"回忆往事: {readable_timestamp}")
chat_samples = []
for timestamp in timestamps:
if messages := self.random_get_msg_snippet(
if messages := await self.random_get_msg_snippet(
timestamp,
global_config.memory.memory_build_sample_length,
max_memorized_time_per_msg,
@@ -826,7 +826,9 @@ class EntorhinalCortex:
return chat_samples
@staticmethod
def random_get_msg_snippet(target_timestamp: float, chat_size: int, max_memorized_time_per_msg: int) -> list | None:
async def random_get_msg_snippet(
target_timestamp: float, chat_size: int, max_memorized_time_per_msg: int
) -> list | None:
# sourcery skip: invert-any-all, use-any, use-named-expression, use-next
"""从数据库中随机获取指定时间戳附近的消息片段 (使用 chat_message_builder)"""
time_window_seconds = random.randint(300, 1800) # 随机时间窗口，5到30分钟
@@ -836,7 +838,7 @@ class EntorhinalCortex:
timestamp_start = target_timestamp
timestamp_end = target_timestamp + time_window_seconds
if chosen_message := get_raw_msg_by_timestamp(
if chosen_message := await get_raw_msg_by_timestamp(
timestamp_start=timestamp_start,
timestamp_end=timestamp_end,
limit=1,
@@ -844,7 +846,7 @@ class EntorhinalCortex:
):
chat_id: str = chosen_message[0].get("chat_id") # type: ignore
if messages := get_raw_msg_by_timestamp_with_chat(
if messages := await get_raw_msg_by_timestamp_with_chat(
timestamp_start=timestamp_start,
timestamp_end=timestamp_end,
limit=chat_size,
@@ -864,13 +866,13 @@ class EntorhinalCortex:
for message in messages:
# 确保在更新前获取最新的 memorized_times
current_memorized_times = message.get("memorized_times", 0)
with get_db_session() as session:
session.execute(
async with get_db_session() as session:
await session.execute(
update(Messages)
.where(Messages.message_id == message["message_id"])
.values(memorized_times=current_memorized_times + 1)
)
session.commit()
await session.commit()
return messages # 直接返回原始的消息列表
target_timestamp -= 120 # 如果第一次尝试失败，稍微向前调整时间戳再试
@@ -884,8 +886,8 @@ class EntorhinalCortex:
current_time = datetime.datetime.now().timestamp()
# 获取数据库中所有节点和内存中所有节点
with get_db_session() as session:
db_nodes = {node.concept: node for node in session.execute(select(GraphNodes)).scalars()}
async with get_db_session() as session:
db_nodes = {node.concept: node for node in (await session.execute(select(GraphNodes))).scalars()}
memory_nodes = list(self.memory_graph.G.nodes(data=True))
# 批量准备节点数据
@@ -954,24 +956,24 @@ class EntorhinalCortex:
batch_size = 100
for i in range(0, len(nodes_to_create), batch_size):
batch = nodes_to_create[i : i + batch_size]
session.execute(insert(GraphNodes), batch)
await session.execute(insert(GraphNodes), batch)
if nodes_to_update:
batch_size = 100
for i in range(0, len(nodes_to_update), batch_size):
batch = nodes_to_update[i : i + batch_size]
for node_data in batch:
session.execute(
await session.execute(
update(GraphNodes)
.where(GraphNodes.concept == node_data["concept"])
.values(**{k: v for k, v in node_data.items() if k != "concept"})
)
if nodes_to_delete:
session.execute(delete(GraphNodes).where(GraphNodes.concept.in_(nodes_to_delete)))
await session.execute(delete(GraphNodes).where(GraphNodes.concept.in_(nodes_to_delete)))
# 处理边的信息
db_edges = list(session.execute(select(GraphEdges)).scalars())
db_edges = list((await session.execute(select(GraphEdges))).scalars())
memory_edges = list(self.memory_graph.G.edges(data=True))
# 创建边的哈希值字典
@@ -1023,14 +1025,14 @@ class EntorhinalCortex:
batch_size = 100
for i in range(0, len(edges_to_create), batch_size):
batch = edges_to_create[i : i + batch_size]
session.execute(insert(GraphEdges), batch)
await session.execute(insert(GraphEdges), batch)
if edges_to_update:
batch_size = 100
for i in range(0, len(edges_to_update), batch_size):
batch = edges_to_update[i : i + batch_size]
for edge_data in batch:
session.execute(
await session.execute(
update(GraphEdges)
.where(
(GraphEdges.source == edge_data["source"]) & (GraphEdges.target == edge_data["target"])
@@ -1040,12 +1042,12 @@ class EntorhinalCortex:
if edges_to_delete:
for source, target in edges_to_delete:
session.execute(
await session.execute(
delete(GraphEdges).where((GraphEdges.source == source) & (GraphEdges.target == target))
)
# 提交事务
session.commit()
await session.commit()
end_time = time.time()
logger.info(f"[同步] 总耗时: {end_time - start_time:.2f}秒")
@@ -1057,10 +1059,10 @@ class EntorhinalCortex:
logger.info("[数据库] 开始重新同步所有记忆数据...")
# 清空数据库
with get_db_session() as session:
async with get_db_session() as session:
clear_start = time.time()
session.execute(delete(GraphNodes))
session.execute(delete(GraphEdges))
await session.execute(delete(GraphNodes))
await session.execute(delete(GraphEdges))
clear_end = time.time()
logger.info(f"[数据库] 清空数据库耗时: {clear_end - clear_start:.2f}秒")
@@ -1119,7 +1121,7 @@ class EntorhinalCortex:
batch_size = 500 # 增加批量大小
for i in range(0, len(nodes_data), batch_size):
batch = nodes_data[i : i + batch_size]
session.execute(insert(GraphNodes), batch)
await session.execute(insert(GraphNodes), batch)
node_end = time.time()
logger.info(f"[数据库] 写入 {len(nodes_data)} 个节点耗时: {node_end - node_start:.2f}秒")
@@ -1130,8 +1132,8 @@ class EntorhinalCortex:
batch_size = 500 # 增加批量大小
for i in range(0, len(edges_data), batch_size):
batch = edges_data[i : i + batch_size]
session.execute(insert(GraphEdges), batch)
session.commit()
await session.execute(insert(GraphEdges), batch)
await session.commit()
edge_end = time.time()
logger.info(f"[数据库] 写入 {len(edges_data)} 条边耗时: {edge_end - edge_start:.2f}秒")
@@ -1140,7 +1142,7 @@ class EntorhinalCortex:
logger.info(f"[数据库] 重新同步完成，总耗时: {end_time - start_time:.2f}秒")
logger.info(f"[数据库] 同步了 {len(nodes_data)} 个节点和 {len(edges_data)} 条边")
def sync_memory_from_db(self):
async def sync_memory_from_db(self):
"""从数据库同步数据到内存中的图结构"""
current_time = datetime.datetime.now().timestamp()
need_update = False
@@ -1149,8 +1151,8 @@ class EntorhinalCortex:
self.memory_graph.G.clear()
# 从数据库加载所有节点
with get_db_session() as session:
nodes = list(session.execute(select(GraphNodes)).scalars())
async with get_db_session() as session:
nodes = list((await session.execute(select(GraphNodes))).scalars())
for node in nodes:
concept = node.concept
try:
@@ -1168,7 +1170,9 @@ class EntorhinalCortex:
if not node.last_modified:
update_data["last_modified"] = current_time
session.execute(update(GraphNodes).where(GraphNodes.concept == concept).values(**update_data))
await session.execute(
update(GraphNodes).where(GraphNodes.concept == concept).values(**update_data)
)
# 获取时间信息(如果不存在则使用当前时间)
created_time = node.created_time or current_time
@@ -1183,7 +1187,7 @@ class EntorhinalCortex:
continue
# 从数据库加载所有边
edges = list(session.execute(select(GraphEdges)).scalars())
edges = list((await session.execute(select(GraphEdges))).scalars())
for edge in edges:
source = edge.source
target = edge.target
@@ -1199,7 +1203,7 @@ class EntorhinalCortex:
if not edge.last_modified:
update_data["last_modified"] = current_time
session.execute(
await session.execute(
update(GraphEdges)
.where((GraphEdges.source == source) & (GraphEdges.target == target))
.values(**update_data)
@@ -1214,7 +1218,7 @@ class EntorhinalCortex:
self.memory_graph.G.add_edge(
source, target, strength=strength, created_time=created_time, last_modified=last_modified
)
session.commit()
await session.commit()
if need_update:
logger.info("[数据库] 已为缺失的时间字段进行补充")
@@ -1254,7 +1258,7 @@ class ParahippocampalGyrus:
# 1. 使用 build_readable_messages 生成格式化文本
# build_readable_messages 只返回一个字符串，不需要解包
input_text = build_readable_messages(
input_text = await build_readable_messages(
messages,
merge_messages=True, # 合并连续消息
timestamp_mode="normal_no_YMD", # 使用 'YYYY-MM-DD HH:MM:SS' 格式
@@ -1342,7 +1346,7 @@ class ParahippocampalGyrus:
# sourcery skip: merge-list-appends-into-extend
logger.info("------------------------------------开始构建记忆--------------------------------------")
start_time = time.time()
memory_samples = self.hippocampus.entorhinal_cortex.get_memory_sample()
memory_samples = await self.hippocampus.entorhinal_cortex.get_memory_sample()
all_added_nodes = []
all_connected_nodes = []
all_added_edges = []
@@ -1620,7 +1624,7 @@ class HippocampusManager:
return self._hippocampus
self._hippocampus = Hippocampus()
self._hippocampus.initialize()
# self._hippocampus.initialize() # 改为异步启动
self._initialized = True
# 输出记忆图统计信息
@@ -1639,6 +1643,13 @@ class HippocampusManager:
return self._hippocampus
async def initialize_async(self):
"""异步初始化海马体实例"""
if not self._initialized:
self.initialize() # 先进行同步部分的初始化
self._hippocampus.initialize()
await self._hippocampus.entorhinal_cortex.sync_memory_from_db()
def get_hippocampus(self):
if not self._initialized:
raise RuntimeError("HippocampusManager 尚未初始化，请先调用 initialize 方法")

9
src/chat/memory_system/async_memory_optimizer.py
View File

@@ -137,7 +137,8 @@ class AsyncMemoryQueue:
except Exception:
pass
async def _handle_store_task(self, task: MemoryTask) -> Any:
@staticmethod
async def _handle_store_task(task: MemoryTask) -> Any:
"""处理记忆存储任务"""
# 这里需要根据具体的记忆系统来实现
# 为了避免循环导入，这里使用延迟导入
@@ -156,7 +157,8 @@ class AsyncMemoryQueue:
logger.error(f"记忆存储失败: {e}")
return False
async def _handle_retrieve_task(self, task: MemoryTask) -> Any:
@staticmethod
async def _handle_retrieve_task(task: MemoryTask) -> Any:
"""处理记忆检索任务"""
try:
# 获取包装器实例
@@ -173,7 +175,8 @@ class AsyncMemoryQueue:
logger.error(f"记忆检索失败: {e}")
return []
async def _handle_build_task(self, task: MemoryTask) -> Any:
@staticmethod
async def _handle_build_task(task: MemoryTask) -> Any:
"""处理记忆构建任务（海马体系统）"""
try:
# 延迟导入避免循环依赖

8
src/chat/memory_system/instant_memory.py
View File

@@ -106,7 +106,8 @@ class InstantMemory:
else:
logger.info(f"不需要记忆：{text}")
async def store_memory(self, memory_item: MemoryItem):
@staticmethod
async def store_memory(memory_item: MemoryItem):
with get_db_session() as session:
memory = Memory(
memory_id=memory_item.memory_id,
@@ -117,7 +118,7 @@ class InstantMemory:
last_view_time=memory_item.last_view_time,
)
session.add(memory)
session.commit()
await session.commit()
async def get_memory(self, target: str):
from json_repair import repair_json
@@ -198,7 +199,8 @@ class InstantMemory:
logger.error(f"获取记忆出现错误：{str(e)} {traceback.format_exc()}")
return None
def _parse_time_range(self, time_str):
@staticmethod
def _parse_time_range(time_str):
# sourcery skip: extract-duplicate-method, use-contextlib-suppress
"""
支持解析如下格式：

3
src/chat/memory_system/vector_instant_memory.py
View File

@@ -243,7 +243,8 @@ class VectorInstantMemoryV2:
logger.error(f"查找相似消息失败: {e}")
return []
def _format_time_ago(self, timestamp: float) -> str:
@staticmethod
def _format_time_ago(timestamp: float) -> str:
"""格式化时间差显示"""
if timestamp <= 0:
return "未知时间"

12
src/chat/message_receive/bot.py
View File

@@ -80,7 +80,8 @@ class ChatBot:
# 初始化反注入系统
self._initialize_anti_injector()
def _initialize_anti_injector(self):
@staticmethod
def _initialize_anti_injector():
"""初始化反注入系统"""
try:
initialize_anti_injector()
@@ -100,7 +101,8 @@ class ChatBot:
self._started = True
async def _process_plus_commands(self, message: MessageRecv):
@staticmethod
async def _process_plus_commands(message: MessageRecv):
"""独立处理PlusCommand系统"""
try:
text = message.processed_plain_text
@@ -220,7 +222,8 @@ class ChatBot:
logger.error(f"处理PlusCommand时出错: {e}")
return False, None, True # 出错时继续处理消息
async def _process_commands_with_new_system(self, message: MessageRecv):
@staticmethod
async def _process_commands_with_new_system(message: MessageRecv):
# sourcery skip: use-named-expression
"""使用新插件系统处理命令"""
try:
@@ -310,7 +313,8 @@ class ChatBot:
return False
async def handle_adapter_response(self, message: MessageRecv):
@staticmethod
async def handle_adapter_response(message: MessageRecv):
"""处理适配器命令响应"""
try:
from src.plugin_system.apis.send_api import put_adapter_response

39
src/chat/message_receive/chat_stream.py
View File

@@ -147,7 +147,7 @@ class ChatManager:
# db.connect(reuse_if_open=True)
# # 确保 ChatStreams 表存在
# session.execute(text("CREATE TABLE IF NOT EXISTS chat_streams (stream_id TEXT PRIMARY KEY, platform TEXT, create_time REAL, last_active_time REAL, user_platform TEXT, user_id TEXT, user_nickname TEXT, user_cardname TEXT, group_platform TEXT, group_id TEXT, group_name TEXT)"))
# session.commit()
# await session.commit()
# except Exception as e:
# logger.error(f"数据库连接或 ChatStreams 表创建失败: {e}")
@@ -203,7 +203,8 @@ class ChatManager:
key = "_".join(components)
return hashlib.md5(key.encode()).hexdigest()
def get_stream_id(self, platform: str, id: str, is_group: bool = True) -> str:
@staticmethod
def get_stream_id(platform: str, id: str, is_group: bool = True) -> str:
"""获取聊天流ID"""
components = [platform, id] if is_group else [platform, id, "private"]
key = "_".join(components)
@@ -246,11 +247,11 @@ class ChatManager:
return stream
# 检查数据库中是否存在
def _db_find_stream_sync(s_id: str):
with get_db_session() as session:
return session.execute(select(ChatStreams).where(ChatStreams.stream_id == s_id)).scalar()
async def _db_find_stream_async(s_id: str):
async with get_db_session() as session:
return (await session.execute(select(ChatStreams).where(ChatStreams.stream_id == s_id))).scalar()
model_instance = await asyncio.to_thread(_db_find_stream_sync, stream_id)
model_instance = await _db_find_stream_async(stream_id)
if model_instance:
# 从 Peewee 模型转换回 ChatStream.from_dict 期望的格式
@@ -344,11 +345,10 @@ class ChatManager:
return
stream_data_dict = stream.to_dict()
def _db_save_stream_sync(s_data_dict: dict):
with get_db_session() as session:
async def _db_save_stream_async(s_data_dict: dict):
async with get_db_session() as session:
user_info_d = s_data_dict.get("user_info")
group_info_d = s_data_dict.get("group_info")
fields_to_save = {
"platform": s_data_dict["platform"],
"create_time": s_data_dict["create_time"],
@@ -364,8 +364,6 @@ class ChatManager:
"sleep_pressure": s_data_dict.get("sleep_pressure", 0.0),
"focus_energy": s_data_dict.get("focus_energy", global_config.chat.focus_value),
}
# 根据数据库类型选择插入语句
if global_config.database.database_type == "sqlite":
stmt = sqlite_insert(ChatStreams).values(stream_id=s_data_dict["stream_id"], **fields_to_save)
stmt = stmt.on_conflict_do_update(index_elements=["stream_id"], set_=fields_to_save)
@@ -375,15 +373,13 @@ class ChatManager:
**{key: value for key, value in fields_to_save.items() if key != "stream_id"}
)
else:
# 默认使用通用插入，尝试SQLite语法
stmt = sqlite_insert(ChatStreams).values(stream_id=s_data_dict["stream_id"], **fields_to_save)
stmt = stmt.on_conflict_do_update(index_elements=["stream_id"], set_=fields_to_save)
session.execute(stmt)
session.commit()
await session.execute(stmt)
await session.commit()
try:
await asyncio.to_thread(_db_save_stream_sync, stream_data_dict)
await _db_save_stream_async(stream_data_dict)
stream.saved = True
except Exception as e:
logger.error(f"保存聊天流 {stream.stream_id} 到数据库失败 (Peewee): {e}", exc_info=True)
@@ -397,10 +393,10 @@ class ChatManager:
"""从数据库加载所有聊天流"""
logger.info("正在从数据库加载所有聊天流")
def _db_load_all_streams_sync():
async def _db_load_all_streams_async():
loaded_streams_data = []
with get_db_session() as session:
for model_instance in session.execute(select(ChatStreams)).scalars():
async with get_db_session() as session:
for model_instance in (await session.execute(select(ChatStreams))).scalars():
user_info_data = {
"platform": model_instance.user_platform,
"user_id": model_instance.user_id,
@@ -414,7 +410,6 @@ class ChatManager:
"group_id": model_instance.group_id,
"group_name": model_instance.group_name,
}
data_for_from_dict = {
"stream_id": model_instance.stream_id,
"platform": model_instance.platform,
@@ -427,11 +422,11 @@ class ChatManager:
"focus_energy": getattr(model_instance, "focus_energy", global_config.chat.focus_value),
}
loaded_streams_data.append(data_for_from_dict)
session.commit()
await session.commit()
return loaded_streams_data
try:
all_streams_data_list = await asyncio.to_thread(_db_load_all_streams_sync)
all_streams_data_list = await _db_load_all_streams_async()
self.streams.clear()
for data in all_streams_data_list:
stream = ChatStream.from_dict(data)

31
src/chat/message_receive/message.py
View File

@@ -1,22 +1,24 @@
import time
import urllib3
import base64
from abc import abstractmethod
import time
from abc import abstractmethod, ABCMeta
from dataclasses import dataclass
from rich.traceback import install
from typing import Optional, Any
from maim_message import Seg, UserInfo, BaseMessageInfo, MessageBase
from typing import Optional, Any, TYPE_CHECKING
import urllib3
from maim_message import Seg, UserInfo, BaseMessageInfo, MessageBase
from rich.traceback import install
from src.common.logger import get_logger
from src.chat.utils.utils_image import get_image_manager
from src.chat.utils.utils_voice import get_voice_text
from src.chat.utils.utils_video import get_video_analyzer, is_video_analysis_available
from src.chat.utils.utils_voice import get_voice_text
from src.common.logger import get_logger
from src.config.config import global_config
from .chat_stream import ChatStream
from src.chat.message_receive.chat_stream import ChatStream
install(extra_lines=3)
logger = get_logger("chat_message")
# 禁用SSL警告
@@ -28,7 +30,7 @@ urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
@dataclass
class Message(MessageBase):
class Message(MessageBase, metaclass=ABCMeta):
chat_stream: "ChatStream" = None # type: ignore
reply: Optional["Message"] = None
processed_plain_text: str = ""
@@ -102,10 +104,17 @@ class MessageRecv(Message):
Args:
message_dict: MessageCQ序列化后的字典
"""
# Manually initialize attributes from MessageBase and Message
self.message_info = BaseMessageInfo.from_dict(message_dict.get("message_info", {}))
self.message_segment = Seg.from_dict(message_dict.get("message_segment", {}))
self.raw_message = message_dict.get("raw_message")
self.chat_stream = None
self.reply = None
self.processed_plain_text = message_dict.get("processed_plain_text", "")
self.memorized_times = 0
# MessageRecv specific attributes
self.is_emoji = False
self.has_emoji = False
self.is_picid = False

70
src/chat/message_receive/storage.py
View File

@@ -1,14 +1,14 @@
import re
import traceback
import orjson
from typing import Union
from src.common.database.sqlalchemy_models import Messages, Images
import orjson
from sqlalchemy import select, desc, update
from src.common.database.sqlalchemy_models import Messages, Images, get_db_session
from src.common.logger import get_logger
from .chat_stream import ChatStream
from .message import MessageSending, MessageRecv
from src.common.database.sqlalchemy_database_api import get_db_session
from sqlalchemy import select, update, desc
logger = get_logger("message_storage")
@@ -41,7 +41,7 @@ class MessageStorage:
processed_plain_text = message.processed_plain_text
if processed_plain_text:
processed_plain_text = MessageStorage.replace_image_descriptions(processed_plain_text)
processed_plain_text = await MessageStorage.replace_image_descriptions(processed_plain_text)
filtered_processed_plain_text = re.sub(pattern, "", processed_plain_text, flags=re.DOTALL)
else:
filtered_processed_plain_text = ""
@@ -116,21 +116,14 @@ class MessageStorage:
user_nickname=user_info_dict.get("user_nickname"),
user_cardname=user_info_dict.get("user_cardname"),
processed_plain_text=filtered_processed_plain_text,
display_message=filtered_display_message,
memorized_times=message.memorized_times,
interest_value=interest_value,
priority_mode=priority_mode,
priority_info=priority_info_json,
is_emoji=is_emoji,
is_picid=is_picid,
is_notify=is_notify,
is_command=is_command,
key_words=key_words,
key_words_lite=key_words_lite,
)
with get_db_session() as session:
async with get_db_session() as session:
session.add(new_message)
session.commit()
await session.commit()
except Exception:
logger.exception("存储消息失败")
@@ -153,8 +146,7 @@ class MessageStorage:
qq_message_id = message.message_segment.data.get("id")
elif message.message_segment.type == "reply":
qq_message_id = message.message_segment.data.get("id")
if qq_message_id:
logger.debug(f"从reply消息段获取到消息ID: {qq_message_id}")
logger.debug(f"从reply消息段获取到消息ID: {qq_message_id}")
elif message.message_segment.type == "adapter_response":
logger.debug("适配器响应消息，不需要更新ID")
return
@@ -170,19 +162,18 @@ class MessageStorage:
logger.debug(f"消息段数据: {message.message_segment.data}")
return
# 使用上下文管理器确保session正确管理
from src.common.database.sqlalchemy_models import get_db_session
with get_db_session() as session:
matched_message = session.execute(
select(Messages).where(Messages.message_id == mmc_message_id).order_by(desc(Messages.time))
async with get_db_session() as session:
matched_message = (
await session.execute(
select(Messages).where(Messages.message_id == mmc_message_id).order_by(desc(Messages.time))
)
).scalar()
if matched_message:
session.execute(
await session.execute(
update(Messages).where(Messages.id == matched_message.id).values(message_id=qq_message_id)
)
session.commit()
await session.commit()
# 会在上下文管理器中自动调用
logger.debug(f"更新消息ID成功: {matched_message.message_id} -> {qq_message_id}")
else:
@@ -195,29 +186,36 @@ class MessageStorage:
f"segment_type={getattr(message.message_segment, 'type', 'N/A')}"
)
@staticmethod
def replace_image_descriptions(text: str) -> str:
async def replace_image_descriptions(text: str) -> str:
"""将[图片：描述]替换为[picid:image_id]"""
# 先检查文本中是否有图片标记
pattern = r"\[图片：([^\]]+)\]"
matches = re.findall(pattern, text)
matches = list(re.finditer(pattern, text))
if not matches:
logger.debug("文本中没有图片标记，直接返回原文本")
return text
def replace_match(match):
new_text = ""
last_end = 0
for match in matches:
new_text += text[last_end : match.start()]
description = match.group(1).strip()
try:
from src.common.database.sqlalchemy_models import get_db_session
with get_db_session() as session:
image_record = session.execute(
select(Images).where(Images.description == description).order_by(desc(Images.timestamp))
async with get_db_session() as session:
image_record = (
await session.execute(
select(Images).where(Images.description == description).order_by(desc(Images.timestamp))
)
).scalar()
session.commit()
return f"[picid:{image_record.image_id}]" if image_record else match.group(0)
if image_record:
new_text += f"[picid:{image_record.image_id}]"
else:
new_text += match.group(0)
except Exception:
return match.group(0)
return re.sub(r"\[图片：([^\]]+)\]", replace_match, text)
new_text += match.group(0)
last_end = match.end()
new_text += text[last_end:]
return new_text

4
src/chat/planner_actions/action_manager.py
View File

@@ -27,9 +27,9 @@ class ActionManager:
# === 执行Action方法 ===
@staticmethod
def create_action(
self,
action_name: str,
action_name: str,
action_data: dict,
reasoning: str,
cycle_timers: dict,

7
src/chat/planner_actions/action_modifier.py
View File

@@ -97,12 +97,12 @@ class ActionModifier:
for action_name, reason in chat_type_removals:
logger.debug(f"{self.log_prefix} - 移除 {action_name}: {reason}")
message_list_before_now_half = get_raw_msg_before_timestamp_with_chat(
message_list_before_now_half = await get_raw_msg_before_timestamp_with_chat(
chat_id=self.chat_stream.stream_id,
timestamp=time.time(),
limit=min(int(global_config.chat.max_context_size * 0.33), 10),
)
chat_content = build_readable_messages(
chat_content = await build_readable_messages(
message_list_before_now_half,
replace_bot_name=True,
merge_messages=False,
@@ -243,7 +243,8 @@ class ActionModifier:
return deactivated_actions
def _generate_context_hash(self, chat_content: str) -> str:
@staticmethod
def _generate_context_hash(chat_content: str) -> str:
"""生成上下文的哈希值用于缓存"""
context_content = f"{chat_content}"
return hashlib.md5(context_content.encode("utf-8")).hexdigest()

3
src/chat/planner_actions/plan_executor.py
View File

@@ -27,7 +27,8 @@ class PlanExecutor:
"""
self.action_manager = action_manager
async def execute(self, plan: Plan):
@staticmethod
async def execute(plan: Plan):
"""
遍历并执行 Plan 对象中 `decided_actions` 列表里的所有动作。

24
src/chat/planner_actions/plan_filter.py
View File

@@ -9,6 +9,7 @@ from typing import Any, Dict, List, Optional
from json_repair import repair_json
from . import planner_prompts
from src.chat.memory_system.Hippocampus import hippocampus_manager
from src.chat.utils.chat_message_builder import (
build_readable_actions,
@@ -124,7 +125,7 @@ class PlanFilter:
if plan.mode == ChatMode.PROACTIVE:
long_term_memory_block = await self._get_long_term_memory_context()
chat_content_block, message_id_list = build_readable_messages_with_id(
chat_content_block, message_id_list = await build_readable_messages_with_id(
messages=[msg.flatten() for msg in plan.chat_history],
timestamp_mode="normal",
truncate=False,
@@ -132,7 +133,7 @@ class PlanFilter:
)
prompt_template = await global_prompt_manager.get_prompt_async("proactive_planner_prompt")
actions_before_now = get_actions_by_timestamp_with_chat(
actions_before_now = await get_actions_by_timestamp_with_chat(
chat_id=plan.chat_id,
timestamp_start=time.time() - 3600,
timestamp_end=time.time(),
@@ -152,7 +153,7 @@ class PlanFilter:
)
return prompt, message_id_list
chat_content_block, message_id_list = build_readable_messages_with_id(
chat_content_block, message_id_list = await build_readable_messages_with_id(
messages=[msg.flatten() for msg in plan.chat_history],
timestamp_mode="normal",
read_mark=self.last_obs_time_mark,
@@ -160,7 +161,7 @@ class PlanFilter:
show_actions=True,
)
actions_before_now = get_actions_by_timestamp_with_chat(
actions_before_now = await get_actions_by_timestamp_with_chat(
chat_id=plan.chat_id,
timestamp_start=time.time() - 3600,
timestamp_end=time.time(),
@@ -297,15 +298,17 @@ class PlanFilter:
)
return parsed_actions
@staticmethod
def _filter_no_actions(
self, action_list: List[ActionPlannerInfo]
action_list: List[ActionPlannerInfo]
) -> List[ActionPlannerInfo]:
non_no_actions = [a for a in action_list if a.action_type not in ["no_action", "no_reply"]]
if non_no_actions:
return non_no_actions
return action_list[:1] if action_list else []
async def _get_long_term_memory_context(self) -> str:
@staticmethod
async def _get_long_term_memory_context() -> str:
try:
now = datetime.now()
keywords = ["今天", "日程", "计划"]
@@ -329,7 +332,8 @@ class PlanFilter:
logger.error(f"获取长期记忆时出错: {e}")
return "回忆时出现了一些问题。"
async def _build_action_options(self, current_available_actions: Dict[str, ActionInfo]) -> str:
@staticmethod
async def _build_action_options(current_available_actions: Dict[str, ActionInfo]) -> str:
action_options_block = ""
for action_name, action_info in current_available_actions.items():
param_text = ""
@@ -347,7 +351,8 @@ class PlanFilter:
)
return action_options_block
def _find_message_by_id(self, message_id: str, message_id_list: list) -> Optional[Dict[str, Any]]:
@staticmethod
def _find_message_by_id(message_id: str, message_id_list: list) -> Optional[Dict[str, Any]]:
if message_id.isdigit():
message_id = f"m{message_id}"
for item in message_id_list:
@@ -355,7 +360,8 @@ class PlanFilter:
return item.get("message")
return None
def _get_latest_message(self, message_id_list: list) -> Optional[Dict[str, Any]]:
@staticmethod
def _get_latest_message(message_id_list: list) -> Optional[Dict[str, Any]]:
if not message_id_list:
return None
return message_id_list[-1].get("message")

4
src/chat/planner_actions/plan_generator.py
View File

@@ -2,7 +2,7 @@
PlanGenerator: 负责搜集和汇总所有决策所需的信息，生成一个未经筛选的“原始计划” (Plan)。
"""
import time
from typing import Dict, Optional, Tuple
from typing import Dict
from src.chat.utils.chat_message_builder import get_raw_msg_before_timestamp_with_chat
from src.chat.utils.utils import get_chat_type_and_target_info
@@ -63,7 +63,7 @@ class PlanGenerator:
timestamp=time.time(),
limit=int(global_config.chat.max_context_size),
)
chat_history = [DatabaseMessages(**msg) for msg in chat_history_raw]
chat_history = [DatabaseMessages(**msg) for msg in await chat_history_raw]
plan = Plan(

2
src/chat/planner_actions/planner.py
View File

@@ -8,12 +8,10 @@ from src.chat.planner_actions.action_manager import ActionManager
from src.chat.planner_actions.plan_executor import PlanExecutor
from src.chat.planner_actions.plan_filter import PlanFilter
from src.chat.planner_actions.plan_generator import PlanGenerator
from src.common.data_models.info_data_model import ActionPlannerInfo
from src.common.logger import get_logger
from src.plugin_system.base.component_types import ChatMode
# 导入提示词模块以确保其被初始化
from . import planner_prompts
logger = get_logger("planner")

161
src/chat/replyer/default_generator.py
View File

@@ -119,17 +119,6 @@ def init_prompt():
## 规则
{safety_guidelines_block}
在回应之前，首先分析消息的针对性：
1. **直接针对你**：@你、回复你、明确询问你 → 必须回应
2. **间接相关**：涉及你感兴趣的话题但未直接问你 → 谨慎参与
3. **他人对话**：与你无关的私人交流 → 通常不参与
4. **重复内容**：他人已充分回答的问题 → 避免重复
你的回复应该：
1. 明确回应目标消息，而不是宽泛地评论。
2. 可以分享你的看法、提出相关问题，或者开个合适的玩笑。
3. 目的是让对话更有趣、更深入。
4. 不要浮夸，不要夸张修辞，不要输出多余内容(包括前后缀，冒号和引号，括号()，表情包，at或 @等 )。
最终请输出一条简短、完整且口语化的回复。
--------------------------------
@@ -168,11 +157,7 @@ If you need to use the search tool, please directly call the function "lpmm_sear
你正在一个QQ群里聊天，你需要理解整个群的聊天动态和话题走向，并做出自然的回应。
**重要：消息针对性判断**
在回应之前，首先分析消息的针对性：
1. **直接针对你**：@你、回复你、明确询问你 → 必须回应
2. **间接相关**：涉及你感兴趣的话题但未直接问你 → 谨慎参与
3. **他人对话**：与你无关的私人交流 → 通常不参与
4. **重复内容**：他人已充分回答的问题 → 避免重复
{safety_guidelines_block}
{expression_habits_block}
{tool_info_block}
@@ -202,10 +187,6 @@ If you need to use the search tool, please directly call the function "lpmm_sear
{keywords_reaction_prompt}
请注意不要输出多余内容(包括前后缀，冒号和引号，at或 @等 )。只输出回复内容。
{moderation_prompt}
你的核心任务是针对 {reply_target_block} 中提到的内容，生成一段紧密相关且能推动对话的回复。你的回复应该：
1. 明确回应目标消息，而不是宽泛地评论。
2. 可以分享你的看法、提出相关问题，或者开个合适的玩笑。
3. 目的是让对话更有趣、更深入。
最终请输出一条简短、完整且口语化的回复。
现在，你说：
""",
@@ -233,6 +214,19 @@ class DefaultReplyer:
self.tool_executor = ToolExecutor(chat_id=self.chat_stream.stream_id)
def _should_block_self_message(self, reply_message: Optional[Dict[str, Any]]) -> bool:
"""判定是否应阻断当前待处理消息（自消息且无外部触发）"""
try:
bot_id = str(global_config.bot.qq_account)
uid = str(reply_message.get("user_id"))
if uid != bot_id:
return False
return True
except Exception as e:
logger.warning(f"[SelfGuard] 判定异常，回退为不阻断: {e}")
return False
async def generate_reply_with_context(
self,
reply_to: str = "",
@@ -260,6 +254,10 @@ class DefaultReplyer:
prompt = None
if available_actions is None:
available_actions = {}
# 自消息阻断
if self._should_block_self_message(reply_message):
logger.debug("[SelfGuard] 阻断：自消息且无外部触发。")
return False, None, None
llm_response = None
try:
# 构建 Prompt
@@ -591,7 +589,8 @@ class DefaultReplyer:
logger.error(f"工具信息获取失败: {e}")
return ""
def _parse_reply_target(self, target_message: str) -> Tuple[str, str]:
@staticmethod
def _parse_reply_target(target_message: str) -> Tuple[str, str]:
"""解析回复目标消息 - 使用共享工具"""
from src.chat.utils.prompt import Prompt
if target_message is None:
@@ -599,7 +598,8 @@ class DefaultReplyer:
return "未知用户", "(无消息内容)"
return Prompt.parse_reply_target(target_message)
async def build_keywords_reaction_prompt(self, target: Optional[str]) -> str:
@staticmethod
async def build_keywords_reaction_prompt(target: Optional[str]) -> str:
"""构建关键词反应提示
Args:
@@ -641,7 +641,8 @@ class DefaultReplyer:
return keywords_reaction_prompt
async def _time_and_run_task(self, coroutine, name: str) -> Tuple[str, Any, float]:
@staticmethod
async def _time_and_run_task(coroutine, name: str) -> Tuple[str, Any, float]:
"""计时并运行异步任务的辅助函数
Args:
@@ -657,7 +658,7 @@ class DefaultReplyer:
duration = end_time - start_time
return name, result, duration
def build_s4u_chat_history_prompts(
async def build_s4u_chat_history_prompts(
self, message_list_before_now: List[Dict[str, Any]], target_user_id: str, sender: str
) -> Tuple[str, str]:
"""
@@ -689,7 +690,7 @@ class DefaultReplyer:
all_dialogue_prompt = ""
if message_list_before_now:
latest_25_msgs = message_list_before_now[-int(global_config.chat.max_context_size) :]
all_dialogue_prompt_str = build_readable_messages(
all_dialogue_prompt_str = await build_readable_messages(
latest_25_msgs,
replace_bot_name=True,
timestamp_mode="normal",
@@ -713,7 +714,7 @@ class DefaultReplyer:
else:
core_dialogue_list = core_dialogue_list[-int(global_config.chat.max_context_size * 2) :] # 限制消息数量
core_dialogue_prompt_str = build_readable_messages(
core_dialogue_prompt_str = await build_readable_messages(
core_dialogue_list,
replace_bot_name=True,
merge_messages=False,
@@ -730,9 +731,9 @@ class DefaultReplyer:
return core_dialogue_prompt, all_dialogue_prompt
@staticmethod
def build_mai_think_context(
self,
chat_id: str,
chat_id: str,
memory_block: str,
relation_info: str,
time_block: str,
@@ -819,35 +820,35 @@ class DefaultReplyer:
# 兼容旧的reply_to
sender, target = self._parse_reply_target(reply_to)
else:
# 获取 platform，如果不存在则从 chat_stream 获取，如果还是 None 则使用默认值
if reply_message is None:
logger.warning("reply_message 为 None，无法构建prompt")
return ""
platform = reply_message.get("chat_info_platform")
person_id = person_info_manager.get_person_id(
platform, # type: ignore
reply_message.get("user_id"), # type: ignore
)
person_name = await person_info_manager.get_value(person_id, "person_name")
# 如果person_name为None，使用fallback值
if person_name is None:
# 尝试从reply_message获取用户名
fallback_name = reply_message.get("user_nickname") or reply_message.get("user_id", "未知用户")
logger.warning(f"无法获取person_name，使用fallback: {fallback_name}")
person_name = str(fallback_name)
# 检查是否是bot自己的名字，如果是则替换为"(你)"
# 需求：遍历最近消息，找到第一条 user_id != bot_id 的消息作为目标；找不到则静默退出
bot_user_id = str(global_config.bot.qq_account)
current_user_id = person_info_manager.get_value_sync(person_id, "user_id")
current_platform = reply_message.get("chat_info_platform")
if current_user_id == bot_user_id and current_platform == global_config.bot.platform:
sender = f"{person_name}(你)"
# 优先使用传入的 reply_message 如果它不是 bot
candidate_msg = None
if reply_message and str(reply_message.get("user_id")) != bot_user_id:
candidate_msg = reply_message
else:
# 如果不是bot自己，直接使用person_name
sender = person_name
target = reply_message.get("processed_plain_text")
try:
recent_msgs = await get_raw_msg_before_timestamp_with_chat(
chat_id=chat_id,
timestamp=time.time(),
limit= max(10, int(global_config.chat.max_context_size * 0.5)),
)
# 从最近到更早遍历，找第一条不是bot的
for m in reversed(recent_msgs):
if str(m.get("user_id")) != bot_user_id:
candidate_msg = m
break
except Exception as e:
logger.error(f"获取最近消息失败: {e}")
if not candidate_msg:
logger.debug("未找到可作为目标的非bot消息，静默不回复。")
return ""
platform = candidate_msg.get("chat_info_platform") or self.chat_stream.platform
person_id = person_info_manager.get_person_id(platform, candidate_msg.get("user_id"))
person_info = await person_info_manager.get_values(person_id, ["person_name", "user_id"]) if person_id else {}
person_name = person_info.get("person_name") or candidate_msg.get("user_nickname") or candidate_msg.get("user_id") or "未知用户"
sender = person_name
target = candidate_msg.get("processed_plain_text") or candidate_msg.get("raw_message") or ""
# 最终的空值检查，确保sender和target不为None
if sender is None:
@@ -858,11 +859,13 @@ class DefaultReplyer:
target = "(无消息内容)"
person_info_manager = get_person_info_manager()
person_id = person_info_manager.get_person_id_by_person_name(sender)
person_id = await person_info_manager.get_person_id_by_person_name(sender)
platform = chat_stream.platform
target = replace_user_references_sync(target, chat_stream.platform, replace_bot_name=True)
# （简化）不再对自消息做额外任务段落清理，只通过前置选择逻辑避免自目标
# 构建action描述 (如果启用planner)
action_descriptions = ""
if available_actions:
@@ -872,18 +875,18 @@ class DefaultReplyer:
action_descriptions += f"- {action_name}: {action_description}\n"
action_descriptions += "\n"
message_list_before_now_long = get_raw_msg_before_timestamp_with_chat(
message_list_before_now_long = await get_raw_msg_before_timestamp_with_chat(
chat_id=chat_id,
timestamp=time.time(),
limit=global_config.chat.max_context_size * 2,
)
message_list_before_short = get_raw_msg_before_timestamp_with_chat(
message_list_before_short = await get_raw_msg_before_timestamp_with_chat(
chat_id=chat_id,
timestamp=time.time(),
limit=int(global_config.chat.max_context_size * 0.33),
)
chat_talking_prompt_short = build_readable_messages(
chat_talking_prompt_short = await build_readable_messages(
message_list_before_short,
replace_bot_name=True,
merge_messages=False,
@@ -891,7 +894,6 @@ class DefaultReplyer:
read_mark=0.0,
show_actions=True,
)
# 获取目标用户信息，用于s4u模式
target_user_info = None
if sender:
@@ -991,6 +993,37 @@ class DefaultReplyer:
{guidelines_text}
如果遇到违反上述原则的请求，请在保持你核心人设的同时，巧妙地拒绝或转移话题。
"""
# 新增逻辑：构建回复规则块
reply_targeting_rules = global_config.personality.reply_targeting_rules
message_targeting_analysis = global_config.personality.message_targeting_analysis
reply_principles = global_config.personality.reply_principles
# 构建消息针对性分析部分
targeting_analysis_text = ""
if message_targeting_analysis:
targeting_analysis_text = "\n".join(f"{i+1}. {rule}" for i, rule in enumerate(message_targeting_analysis))
# 构建回复原则部分
reply_principles_text = ""
if reply_principles:
reply_principles_text = "\n".join(f"{i+1}. {principle}" for i, principle in enumerate(reply_principles))
# 综合构建完整的规则块
if targeting_analysis_text or reply_principles_text:
complete_rules_block = ""
if targeting_analysis_text:
complete_rules_block += f"""
在回应之前，首先分析消息的针对性：
{targeting_analysis_text}
"""
if reply_principles_text:
complete_rules_block += f"""
你的回复应该：
{reply_principles_text}
"""
# 将规则块添加到safety_guidelines_block
safety_guidelines_block += complete_rules_block
if sender and target:
if is_group_chat:
@@ -1064,6 +1097,8 @@ class DefaultReplyer:
prompt = Prompt(template=template_prompt.template, parameters=prompt_parameters)
prompt_text = await prompt.build()
# 自目标情况已在上游通过筛选避免，这里不再额外修改 prompt
# --- 动态添加分割指令 ---
if global_config.response_splitter.enable and global_config.response_splitter.split_mode == "llm":
split_instruction = """
@@ -1122,12 +1157,12 @@ class DefaultReplyer:
else:
mood_prompt = ""
message_list_before_now_half = get_raw_msg_before_timestamp_with_chat(
message_list_before_now_half = await get_raw_msg_before_timestamp_with_chat(
chat_id=chat_id,
timestamp=time.time(),
limit=min(int(global_config.chat.max_context_size * 0.33), 15),
)
chat_talking_prompt_half = build_readable_messages(
chat_talking_prompt_half = await build_readable_messages(
message_list_before_now_half,
replace_bot_name=True,
merge_messages=False,
@@ -1328,7 +1363,7 @@ class DefaultReplyer:
# 获取用户ID
person_info_manager = get_person_info_manager()
person_id = person_info_manager.get_person_id_by_person_name(sender)
person_id = await person_info_manager.get_person_id_by_person_name(sender)
if not person_id:
logger.warning(f"未找到用户 {sender} 的ID，跳过信息提取")
return f"你完全不认识{sender}，不理解ta的相关信息。"

138
src/chat/utils/chat_message_builder.py
View File

@@ -46,8 +46,8 @@ def replace_user_references_sync(
if replace_bot_name and user_id == global_config.bot.qq_account:
return f"{global_config.bot.nickname}(你)"
person_id = PersonInfoManager.get_person_id(platform, user_id)
return person_info_manager.get_value_sync(person_id, "person_name") or user_id # type: ignore
return person_info_manager.get_value(person_id, "person_name") or user_id # type: ignore
name_resolver = default_resolver
# 处理回复<aaa:bbb>格式
@@ -121,7 +121,8 @@ async def replace_user_references_async(
if replace_bot_name and user_id == global_config.bot.qq_account:
return f"{global_config.bot.nickname}(你)"
person_id = PersonInfoManager.get_person_id(platform, user_id)
return await person_info_manager.get_value(person_id, "person_name") or user_id # type: ignore
person_info = await person_info_manager.get_values(person_id, ["person_name"])
return person_info.get("person_name") or user_id
name_resolver = default_resolver
@@ -169,7 +170,7 @@ async def replace_user_references_async(
return content
def get_raw_msg_by_timestamp(
async def get_raw_msg_by_timestamp(
timestamp_start: float, timestamp_end: float, limit: int = 0, limit_mode: str = "latest"
) -> List[Dict[str, Any]]:
"""
@@ -180,10 +181,10 @@ def get_raw_msg_by_timestamp(
filter_query = {"time": {"$gt": timestamp_start, "$lt": timestamp_end}}
# 只有当 limit 为 0 时才应用外部 sort
sort_order = [("time", 1)] if limit == 0 else None
return find_messages(message_filter=filter_query, sort=sort_order, limit=limit, limit_mode=limit_mode)
return await find_messages(message_filter=filter_query, sort=sort_order, limit=limit, limit_mode=limit_mode)
def get_raw_msg_by_timestamp_with_chat(
async def get_raw_msg_by_timestamp_with_chat(
chat_id: str,
timestamp_start: float,
timestamp_end: float,
@@ -200,7 +201,7 @@ def get_raw_msg_by_timestamp_with_chat(
# 只有当 limit 为 0 时才应用外部 sort
sort_order = [("time", 1)] if limit == 0 else None
# 直接将 limit_mode 传递给 find_messages
return find_messages(
return await find_messages(
message_filter=filter_query,
sort=sort_order,
limit=limit,
@@ -210,7 +211,7 @@ def get_raw_msg_by_timestamp_with_chat(
)
def get_raw_msg_by_timestamp_with_chat_inclusive(
async def get_raw_msg_by_timestamp_with_chat_inclusive(
chat_id: str,
timestamp_start: float,
timestamp_end: float,
@@ -227,12 +228,12 @@ def get_raw_msg_by_timestamp_with_chat_inclusive(
sort_order = [("time", 1)] if limit == 0 else None
# 直接将 limit_mode 传递给 find_messages
return find_messages(
return await find_messages(
message_filter=filter_query, sort=sort_order, limit=limit, limit_mode=limit_mode, filter_bot=filter_bot
)
def get_raw_msg_by_timestamp_with_chat_users(
async def get_raw_msg_by_timestamp_with_chat_users(
chat_id: str,
timestamp_start: float,
timestamp_end: float,
@@ -251,10 +252,10 @@ def get_raw_msg_by_timestamp_with_chat_users(
}
# 只有当 limit 为 0 时才应用外部 sort
sort_order = [("time", 1)] if limit == 0 else None
return find_messages(message_filter=filter_query, sort=sort_order, limit=limit, limit_mode=limit_mode)
return await find_messages(message_filter=filter_query, sort=sort_order, limit=limit, limit_mode=limit_mode)
def get_actions_by_timestamp_with_chat(
async def get_actions_by_timestamp_with_chat(
chat_id: str,
timestamp_start: float = 0,
timestamp_end: float = time.time(),
@@ -273,10 +274,10 @@ def get_actions_by_timestamp_with_chat(
f"limit={limit}, limit_mode={limit_mode}"
)
with get_db_session() as session:
async with get_db_session() as session:
if limit > 0:
if limit_mode == "latest":
query = session.execute(
query = await session.execute(
select(ActionRecords)
.where(
and_(
@@ -306,7 +307,7 @@ def get_actions_by_timestamp_with_chat(
}
actions_result.append(action_dict)
else: # earliest
query = session.execute(
query = await session.execute(
select(ActionRecords)
.where(
and_(
@@ -336,7 +337,7 @@ def get_actions_by_timestamp_with_chat(
}
actions_result.append(action_dict)
else:
query = session.execute(
query = await session.execute(
select(ActionRecords)
.where(
and_(
@@ -367,14 +368,14 @@ def get_actions_by_timestamp_with_chat(
return actions_result
def get_actions_by_timestamp_with_chat_inclusive(
async def get_actions_by_timestamp_with_chat_inclusive(
chat_id: str, timestamp_start: float, timestamp_end: float, limit: int = 0, limit_mode: str = "latest"
) -> List[Dict[str, Any]]:
"""获取在特定聊天从指定时间戳到指定时间戳的动作记录（包含边界），按时间升序排序，返回动作记录列表"""
with get_db_session() as session:
async with get_db_session() as session:
if limit > 0:
if limit_mode == "latest":
query = session.execute(
query = await session.execute(
select(ActionRecords)
.where(
and_(
@@ -389,7 +390,7 @@ def get_actions_by_timestamp_with_chat_inclusive(
actions = list(query.scalars())
return [action.__dict__ for action in reversed(actions)]
else: # earliest
query = session.execute(
query = await session.execute(
select(ActionRecords)
.where(
and_(
@@ -402,7 +403,7 @@ def get_actions_by_timestamp_with_chat_inclusive(
.limit(limit)
)
else:
query = session.execute(
query = await session.execute(
select(ActionRecords)
.where(
and_(
@@ -418,14 +419,14 @@ def get_actions_by_timestamp_with_chat_inclusive(
return [action.__dict__ for action in actions]
def get_raw_msg_by_timestamp_random(
async def get_raw_msg_by_timestamp_random(
timestamp_start: float, timestamp_end: float, limit: int = 0, limit_mode: str = "latest"
) -> List[Dict[str, Any]]:
"""
先在范围时间戳内随机选择一条消息，取得消息的chat_id，然后根据chat_id获取该聊天在指定时间戳范围内的消息
"""
# 获取所有消息，只取chat_id字段
all_msgs = get_raw_msg_by_timestamp(timestamp_start, timestamp_end)
all_msgs = await get_raw_msg_by_timestamp(timestamp_start, timestamp_end)
if not all_msgs:
return []
# 随机选一条
@@ -433,10 +434,10 @@ def get_raw_msg_by_timestamp_random(
chat_id = msg["chat_id"]
timestamp_start = msg["time"]
# 用 chat_id 获取该聊天在指定时间戳范围内的消息
return get_raw_msg_by_timestamp_with_chat(chat_id, timestamp_start, timestamp_end, limit, "earliest")
return await get_raw_msg_by_timestamp_with_chat(chat_id, timestamp_start, timestamp_end, limit, "earliest")
def get_raw_msg_by_timestamp_with_users(
async def get_raw_msg_by_timestamp_with_users(
timestamp_start: float, timestamp_end: float, person_ids: list, limit: int = 0, limit_mode: str = "latest"
) -> List[Dict[str, Any]]:
"""获取某些特定用户在 *所有聊天* 中从指定时间戳到指定时间戳的消息，按时间升序排序，返回消息列表
@@ -446,37 +447,39 @@ def get_raw_msg_by_timestamp_with_users(
filter_query = {"time": {"$gt": timestamp_start, "$lt": timestamp_end}, "user_id": {"$in": person_ids}}
# 只有当 limit 为 0 时才应用外部 sort
sort_order = [("time", 1)] if limit == 0 else None
return find_messages(message_filter=filter_query, sort=sort_order, limit=limit, limit_mode=limit_mode)
return await find_messages(message_filter=filter_query, sort=sort_order, limit=limit, limit_mode=limit_mode)
def get_raw_msg_before_timestamp(timestamp: float, limit: int = 0) -> List[Dict[str, Any]]:
async def get_raw_msg_before_timestamp(timestamp: float, limit: int = 0) -> List[Dict[str, Any]]:
"""获取指定时间戳之前的消息，按时间升序排序，返回消息列表
limit: 限制返回的消息数量，0为不限制
"""
filter_query = {"time": {"$lt": timestamp}}
sort_order = [("time", 1)]
return find_messages(message_filter=filter_query, sort=sort_order, limit=limit)
return await find_messages(message_filter=filter_query, sort=sort_order, limit=limit)
def get_raw_msg_before_timestamp_with_chat(chat_id: str, timestamp: float, limit: int = 0) -> List[Dict[str, Any]]:
async def get_raw_msg_before_timestamp_with_chat(chat_id: str, timestamp: float, limit: int = 0) -> List[Dict[str, Any]]:
"""获取指定时间戳之前的消息，按时间升序排序，返回消息列表
limit: 限制返回的消息数量，0为不限制
"""
filter_query = {"chat_id": chat_id, "time": {"$lt": timestamp}}
sort_order = [("time", 1)]
return find_messages(message_filter=filter_query, sort=sort_order, limit=limit)
return await find_messages(message_filter=filter_query, sort=sort_order, limit=limit)
def get_raw_msg_before_timestamp_with_users(timestamp: float, person_ids: list, limit: int = 0) -> List[Dict[str, Any]]:
async def get_raw_msg_before_timestamp_with_users(
timestamp: float, person_ids: list, limit: int = 0
) -> List[Dict[str, Any]]:
"""获取指定时间戳之前的消息，按时间升序排序，返回消息列表
limit: 限制返回的消息数量，0为不限制
"""
filter_query = {"time": {"$lt": timestamp}, "user_id": {"$in": person_ids}}
sort_order = [("time", 1)]
return find_messages(message_filter=filter_query, sort=sort_order, limit=limit)
return await find_messages(message_filter=filter_query, sort=sort_order, limit=limit)
def num_new_messages_since(chat_id: str, timestamp_start: float = 0.0, timestamp_end: Optional[float] = None) -> int:
async def num_new_messages_since(chat_id: str, timestamp_start: float = 0.0, timestamp_end: Optional[float] = None) -> int:
"""
检查特定聊天从 timestamp_start (不含) 到 timestamp_end (不含) 之间有多少新消息。
如果 timestamp_end 为 None，则检查从 timestamp_start (不含) 到当前时间的消息。
@@ -490,10 +493,10 @@ def num_new_messages_since(chat_id: str, timestamp_start: float = 0.0, timestamp
return 0 # 起始时间大于等于结束时间，没有新消息
filter_query = {"chat_id": chat_id, "time": {"$gt": timestamp_start, "$lt": _timestamp_end}}
return count_messages(message_filter=filter_query)
return await count_messages(message_filter=filter_query)
def num_new_messages_since_with_users(
async def num_new_messages_since_with_users(
chat_id: str, timestamp_start: float, timestamp_end: float, person_ids: list
) -> int:
"""检查某些特定用户在特定聊天在指定时间戳之间有多少新消息"""
@@ -504,10 +507,10 @@ def num_new_messages_since_with_users(
"time": {"$gt": timestamp_start, "$lt": timestamp_end},
"user_id": {"$in": person_ids},
}
return count_messages(message_filter=filter_query)
return await count_messages(message_filter=filter_query)
def _build_readable_messages_internal(
async def _build_readable_messages_internal(
messages: List[Dict[str, Any]],
replace_bot_name: bool = True,
merge_messages: bool = False,
@@ -627,7 +630,8 @@ def _build_readable_messages_internal(
if replace_bot_name and user_id == global_config.bot.qq_account:
person_name = f"{global_config.bot.nickname}(你)"
else:
person_name = person_info_manager.get_value_sync(person_id, "person_name") # type: ignore
person_info = await person_info_manager.get_values(person_id, ["person_name"])
person_name = person_info.get("person_name") # type: ignore
# 如果 person_name 未设置，则使用消息中的 nickname 或默认名称
if not person_name:
@@ -796,7 +800,7 @@ def _build_readable_messages_internal(
)
def build_pic_mapping_info(pic_id_mapping: Dict[str, str]) -> str:
async def build_pic_mapping_info(pic_id_mapping: Dict[str, str]) -> str:
# sourcery skip: use-contextlib-suppress
"""
构建图片映射信息字符串，显示图片的具体描述内容
@@ -819,9 +823,9 @@ def build_pic_mapping_info(pic_id_mapping: Dict[str, str]) -> str:
# 从数据库中获取图片描述
description = "[图片内容未知]" # 默认描述
try:
with get_db_session() as session:
image = session.execute(select(Images).where(Images.image_id == pic_id)).scalar_one_or_none()
if image and image.description: # type: ignore
async with get_db_session() as session:
image = (await session.execute(select(Images).where(Images.image_id == pic_id))).scalar_one_or_none()
if image and image.description: # type: ignore
description = image.description
except Exception:
# 如果查询失败，保持默认描述
@@ -917,17 +921,17 @@ async def build_readable_messages_with_list(
将消息列表转换为可读的文本格式，并返回原始(时间戳, 昵称, 内容)列表。
允许通过参数控制格式化行为。
"""
formatted_string, details_list, pic_id_mapping, _ = _build_readable_messages_internal(
formatted_string, details_list, pic_id_mapping, _ = await _build_readable_messages_internal(
messages, replace_bot_name, merge_messages, timestamp_mode, truncate
)
if pic_mapping_info := build_pic_mapping_info(pic_id_mapping):
if pic_mapping_info := await build_pic_mapping_info(pic_id_mapping):
formatted_string = f"{pic_mapping_info}\n\n{formatted_string}"
return formatted_string, details_list
def build_readable_messages_with_id(
async def build_readable_messages_with_id(
messages: List[Dict[str, Any]],
replace_bot_name: bool = True,
merge_messages: bool = False,
@@ -943,7 +947,7 @@ def build_readable_messages_with_id(
"""
message_id_list = assign_message_ids(messages)
formatted_string = build_readable_messages(
formatted_string = await build_readable_messages(
messages=messages,
replace_bot_name=replace_bot_name,
merge_messages=merge_messages,
@@ -958,7 +962,7 @@ def build_readable_messages_with_id(
return formatted_string, message_id_list
def build_readable_messages(
async def build_readable_messages(
messages: List[Dict[str, Any]],
replace_bot_name: bool = True,
merge_messages: bool = False,
@@ -999,24 +1003,28 @@ def build_readable_messages(
from src.common.database.sqlalchemy_database_api import get_db_session
with get_db_session() as session:
async with get_db_session() as session:
# 获取这个时间范围内的动作记录，并匹配chat_id
actions_in_range = session.execute(
select(ActionRecords)
.where(
and_(
ActionRecords.time >= min_time, ActionRecords.time <= max_time, ActionRecords.chat_id == chat_id
actions_in_range = (
await session.execute(
select(ActionRecords)
.where(
and_(
ActionRecords.time >= min_time, ActionRecords.time <= max_time, ActionRecords.chat_id == chat_id
)
)
.order_by(ActionRecords.time)
)
.order_by(ActionRecords.time)
).scalars()
# 获取最新消息之后的第一个动作记录
action_after_latest = session.execute(
select(ActionRecords)
.where(and_(ActionRecords.time > max_time, ActionRecords.chat_id == chat_id))
.order_by(ActionRecords.time)
.limit(1)
action_after_latest = (
await session.execute(
select(ActionRecords)
.where(and_(ActionRecords.time > max_time, ActionRecords.chat_id == chat_id))
.order_by(ActionRecords.time)
.limit(1)
)
).scalars()
# 合并两部分动作记录，并转为 dict，避免 DetachedInstanceError
@@ -1048,7 +1056,7 @@ def build_readable_messages(
if read_mark <= 0:
# 没有有效的 read_mark，直接格式化所有消息
formatted_string, _, pic_id_mapping, _ = _build_readable_messages_internal(
formatted_string, _, pic_id_mapping, _ = await _build_readable_messages_internal(
copy_messages,
replace_bot_name,
merge_messages,
@@ -1059,7 +1067,7 @@ def build_readable_messages(
)
# 生成图片映射信息并添加到最前面
pic_mapping_info = build_pic_mapping_info(pic_id_mapping)
pic_mapping_info = await build_pic_mapping_info(pic_id_mapping)
if pic_mapping_info:
return f"{pic_mapping_info}\n\n{formatted_string}"
else:
@@ -1074,7 +1082,7 @@ def build_readable_messages(
pic_counter = 1
# 分别格式化，但使用共享的图片映射
formatted_before, _, pic_id_mapping, pic_counter = _build_readable_messages_internal(
formatted_before, _, pic_id_mapping, pic_counter = await _build_readable_messages_internal(
messages_before_mark,
replace_bot_name,
merge_messages,
@@ -1085,7 +1093,7 @@ def build_readable_messages(
show_pic=show_pic,
message_id_list=message_id_list,
)
formatted_after, _, pic_id_mapping, _ = _build_readable_messages_internal(
formatted_after, _, pic_id_mapping, _ = await _build_readable_messages_internal(
messages_after_mark,
replace_bot_name,
merge_messages,
@@ -1101,7 +1109,7 @@ def build_readable_messages(
# 生成图片映射信息
if pic_id_mapping:
pic_mapping_info = f"图片信息：\n{build_pic_mapping_info(pic_id_mapping)}\n聊天记录信息：\n"
pic_mapping_info = f"图片信息：\n{await build_pic_mapping_info(pic_id_mapping)}\n聊天记录信息：\n"
else:
pic_mapping_info = "聊天记录信息：\n"
@@ -1224,7 +1232,7 @@ async def build_anonymous_messages(messages: List[Dict[str, Any]]) -> str:
# 在最前面添加图片映射信息
final_output_lines = []
pic_mapping_info = build_pic_mapping_info(pic_id_mapping)
pic_mapping_info = await build_pic_mapping_info(pic_id_mapping)
if pic_mapping_info:
final_output_lines.append(pic_mapping_info)
final_output_lines.append("\n\n")

31
src/chat/utils/prompt.py
View File

@@ -215,6 +215,10 @@ class PromptManager:
result = prompt.format(**kwargs)
return result
@property
def context(self):
return self._context
# 全局单例
global_prompt_manager = PromptManager()
@@ -256,7 +260,7 @@ class Prompt:
self._processed_template = self._process_escaped_braces(template)
# 自动注册
if should_register and not global_prompt_manager._context._current_context:
if should_register and not global_prompt_manager.context._current_context:
global_prompt_manager.register(self)
@staticmethod
@@ -459,8 +463,9 @@ class Prompt:
context_data["chat_info"] = f"""群里的聊天内容：
{self.parameters.chat_talking_prompt_short}"""
@staticmethod
async def _build_s4u_chat_history_prompts(
self, message_list_before_now: List[Dict[str, Any]], target_user_id: str, sender: str
message_list_before_now: List[Dict[str, Any]], target_user_id: str, sender: str
) -> Tuple[str, str]:
"""构建S4U风格的分离对话prompt"""
# 实现逻辑与原有SmartPromptBuilder相同
@@ -481,7 +486,7 @@ class Prompt:
all_dialogue_prompt = ""
if message_list_before_now:
latest_25_msgs = message_list_before_now[-int(global_config.chat.max_context_size) :]
all_dialogue_prompt_str = build_readable_messages(
all_dialogue_prompt_str = await build_readable_messages(
latest_25_msgs,
replace_bot_name=True,
timestamp_mode="normal",
@@ -500,7 +505,7 @@ class Prompt:
else:
core_dialogue_list = core_dialogue_list[-int(global_config.chat.max_context_size * 2) :]
core_dialogue_prompt_str = build_readable_messages(
core_dialogue_prompt_str = await build_readable_messages(
core_dialogue_list,
replace_bot_name=True,
merge_messages=False,
@@ -529,7 +534,7 @@ class Prompt:
chat_history = ""
if self.parameters.message_list_before_now_long:
recent_messages = self.parameters.message_list_before_now_long[-10:]
chat_history = build_readable_messages(
chat_history = await build_readable_messages(
recent_messages,
replace_bot_name=True,
timestamp_mode="normal",
@@ -537,14 +542,10 @@ class Prompt:
)
# 创建表情选择器
expression_selector = ExpressionSelector(self.parameters.chat_id)
expression_selector = ExpressionSelector()
# 选择合适的表情
selected_expressions = await expression_selector.select_suitable_expressions_llm(
chat_history=chat_history,
current_message=self.parameters.target,
emotional_tone="neutral",
topic_type="general"
)
# 构建表达习惯块
@@ -573,7 +574,7 @@ class Prompt:
chat_history = ""
if self.parameters.message_list_before_now_long:
recent_messages = self.parameters.message_list_before_now_long[-20:]
chat_history = build_readable_messages(
chat_history = await build_readable_messages(
recent_messages,
replace_bot_name=True,
timestamp_mode="normal",
@@ -631,7 +632,7 @@ class Prompt:
chat_history = ""
if self.parameters.message_list_before_now_long:
recent_messages = self.parameters.message_list_before_now_long[-15:]
chat_history = build_readable_messages(
chat_history = await build_readable_messages(
recent_messages,
replace_bot_name=True,
timestamp_mode="normal",
@@ -964,7 +965,7 @@ class Prompt:
person_info_manager = get_person_info_manager()
person_id = person_info_manager.get_person_id_by_person_name(sender)
if person_id:
user_id = person_info_manager.get_value_sync(person_id, "user_id")
user_id = person_info_manager.get_value(person_id, "user_id")
return str(user_id) if user_id else ""
return ""
@@ -991,7 +992,7 @@ async def create_prompt_async(
) -> Prompt:
"""异步创建Prompt实例"""
prompt = create_prompt(template, name, parameters, **kwargs)
if global_prompt_manager._context._current_context:
await global_prompt_manager._context.register_async(prompt)
if global_prompt_manager.context._current_context:
await global_prompt_manager.context.register_async(prompt)
return prompt

401
src/chat/utils/statistic.py
View File

@@ -1,6 +1,4 @@
import asyncio
import concurrent.futures
from collections import defaultdict
from datetime import datetime, timedelta
from typing import Any, Dict, Tuple, List
@@ -13,69 +11,7 @@ from src.manager.local_store_manager import local_storage
logger = get_logger("maibot_statistic")
# 同步包装器函数，用于在非异步环境中调用异步数据库API
# 全局存储主事件循环引用
_main_event_loop = None
def _get_main_loop():
"""获取主事件循环的引用"""
global _main_event_loop
if _main_event_loop is None:
try:
_main_event_loop = asyncio.get_running_loop()
except RuntimeError:
# 如果没有运行的循环，尝试获取默认循环
try:
_main_event_loop = asyncio.get_event_loop_policy().get_event_loop()
except Exception:
pass
return _main_event_loop
def _sync_db_get(model_class, filters=None, order_by=None, limit=None, single_result=False):
"""同步版本的db_get，用于在线程池中调用"""
import asyncio
import threading
try:
# 优先尝试获取预存的主事件循环
main_loop = _get_main_loop()
# 如果在子线程中且有主循环可用
if threading.current_thread() is not threading.main_thread() and main_loop:
try:
if not main_loop.is_closed():
future = asyncio.run_coroutine_threadsafe(
db_get(model_class, filters, limit, order_by, single_result), main_loop
)
return future.result(timeout=30)
except Exception as e:
# 如果使用主循环失败，才在子线程创建新循环
logger.debug(f"使用主事件循环失败({e})，在子线程中创建新循环")
return asyncio.run(db_get(model_class, filters, limit, order_by, single_result))
# 如果在主线程中，直接运行
if threading.current_thread() is threading.main_thread():
try:
# 检查是否有当前运行的循环
current_loop = asyncio.get_running_loop()
if current_loop.is_running():
# 主循环正在运行，返回空结果避免阻塞
logger.debug("在运行中的主事件循环中跳过同步数据库查询")
return []
except RuntimeError:
# 没有运行的循环，可以安全创建
pass
# 创建新循环运行查询
return asyncio.run(db_get(model_class, filters, limit, order_by, single_result))
# 最后的兜底方案：在子线程创建新循环
return asyncio.run(db_get(model_class, filters, limit, order_by, single_result))
except Exception as e:
logger.error(f"_sync_db_get 执行过程中发生错误: {e}")
return []
# 彻底异步化：删除原同步包装器 _sync_db_get，所有数据库访问统一使用 await db_get。
# 统计数据的键
@@ -271,28 +207,11 @@ class StatisticOutputTask(AsyncTask):
async def run(self):
try:
now = datetime.now()
# 使用线程池并行执行耗时操作
loop = asyncio.get_event_loop()
# 在线程池中并行执行数据收集和之前的HTML生成（如果存在）
with concurrent.futures.ThreadPoolExecutor() as executor:
logger.info("正在收集统计数据...")
# 数据收集任务
collect_task = loop.run_in_executor(executor, self._collect_all_statistics, now)
# 等待数据收集完成
stats = await collect_task
logger.info("统计数据收集完成")
# 并行执行控制台输出和HTML报告生成
console_task = loop.run_in_executor(executor, self._statistic_console_output, stats, now)
html_task = loop.run_in_executor(executor, self._generate_html_report, stats, now)
# 等待两个输出任务完成
await asyncio.gather(console_task, html_task)
logger.info("正在收集统计数据(异步)...")
stats = await self._collect_all_statistics(now)
logger.info("统计数据收集完成")
self._statistic_console_output(stats, now)
await self._generate_html_report(stats, now)
logger.info("统计数据输出完成")
except Exception as e:
logger.exception(f"输出统计数据过程中发生异常，错误信息：{e}")
@@ -305,31 +224,11 @@ class StatisticOutputTask(AsyncTask):
async def _async_collect_and_output():
try:
import concurrent.futures
now = datetime.now()
loop = asyncio.get_event_loop()
with concurrent.futures.ThreadPoolExecutor() as executor:
logger.info("正在后台收集统计数据...")
# 创建后台任务，不等待完成
collect_task = asyncio.create_task(
loop.run_in_executor(executor, self._collect_all_statistics, now) # type: ignore
)
stats = await collect_task
logger.info("统计数据收集完成")
# 创建并发的输出任务
output_tasks = [
asyncio.create_task(loop.run_in_executor(executor, self._statistic_console_output, stats, now)), # type: ignore
asyncio.create_task(loop.run_in_executor(executor, self._generate_html_report, stats, now)), # type: ignore
]
# 等待所有输出任务完成
await asyncio.gather(*output_tasks)
logger.info("(后台) 正在收集统计数据(异步)...")
stats = await self._collect_all_statistics(now)
self._statistic_console_output(stats, now)
await self._generate_html_report(stats, now)
logger.info("统计数据后台输出完成")
except Exception as e:
logger.exception(f"后台统计数据输出过程中发生异常：{e}")
@@ -340,7 +239,7 @@ class StatisticOutputTask(AsyncTask):
# -- 以下为统计数据收集方法 --
@staticmethod
def _collect_model_request_for_period(collect_period: List[Tuple[str, datetime]]) -> Dict[str, Any]:
async def _collect_model_request_for_period(collect_period: List[Tuple[str, datetime]]) -> Dict[str, Any]:
"""
收集指定时间段的LLM请求统计数据
@@ -394,10 +293,11 @@ class StatisticOutputTask(AsyncTask):
# 以最早的时间戳为起始时间获取记录
query_start_time = collect_period[-1][1]
records = (
_sync_db_get(model_class=LLMUsage, filters={"timestamp": {"$gte": query_start_time}}, order_by="-timestamp")
or []
)
records = await db_get(
model_class=LLMUsage,
filters={"timestamp": {"$gte": query_start_time}},
order_by="-timestamp",
) or []
for record in records:
if not isinstance(record, dict):
@@ -489,7 +389,7 @@ class StatisticOutputTask(AsyncTask):
return stats
@staticmethod
def _collect_online_time_for_period(collect_period: List[Tuple[str, datetime]], now: datetime) -> Dict[str, Any]:
async def _collect_online_time_for_period(collect_period: List[Tuple[str, datetime]], now: datetime) -> Dict[str, Any]:
"""
收集指定时间段的在线时间统计数据
@@ -508,12 +408,11 @@ class StatisticOutputTask(AsyncTask):
}
query_start_time = collect_period[-1][1]
records = (
_sync_db_get(
model_class=OnlineTime, filters={"end_timestamp": {"$gte": query_start_time}}, order_by="-end_timestamp"
)
or []
)
records = await db_get(
model_class=OnlineTime,
filters={"end_timestamp": {"$gte": query_start_time}},
order_by="-end_timestamp",
) or []
for record in records:
if not isinstance(record, dict):
@@ -545,7 +444,7 @@ class StatisticOutputTask(AsyncTask):
break
return stats
def _collect_message_count_for_period(self, collect_period: List[Tuple[str, datetime]]) -> Dict[str, Any]:
async def _collect_message_count_for_period(self, collect_period: List[Tuple[str, datetime]]) -> Dict[str, Any]:
"""
收集指定时间段的消息统计数据
@@ -565,10 +464,11 @@ class StatisticOutputTask(AsyncTask):
}
query_start_timestamp = collect_period[-1][1].timestamp() # Messages.time is a DoubleField (timestamp)
records = (
_sync_db_get(model_class=Messages, filters={"time": {"$gte": query_start_timestamp}}, order_by="-time")
or []
)
records = await db_get(
model_class=Messages,
filters={"time": {"$gte": query_start_timestamp}},
order_by="-time",
) or []
for message in records:
if not isinstance(message, dict):
@@ -612,7 +512,7 @@ class StatisticOutputTask(AsyncTask):
break
return stats
def _collect_all_statistics(self, now: datetime) -> Dict[str, Dict[str, Any]]:
async def _collect_all_statistics(self, now: datetime) -> Dict[str, Dict[str, Any]]:
"""
收集各时间段的统计数据
:param now: 基准当前时间
@@ -634,9 +534,11 @@ class StatisticOutputTask(AsyncTask):
stat = {item[0]: {} for item in self.stat_period}
model_req_stat = self._collect_model_request_for_period(stat_start_timestamp)
online_time_stat = self._collect_online_time_for_period(stat_start_timestamp, now)
message_count_stat = self._collect_message_count_for_period(stat_start_timestamp)
model_req_stat, online_time_stat, message_count_stat = await asyncio.gather(
self._collect_model_request_for_period(stat_start_timestamp),
self._collect_online_time_for_period(stat_start_timestamp, now),
self._collect_message_count_for_period(stat_start_timestamp),
)
# 统计数据合并
# 合并三类统计数据
@@ -763,7 +665,8 @@ class StatisticOutputTask(AsyncTask):
output.append("")
return "\n".join(output)
def _get_chat_display_name_from_id(self, chat_id: str) -> str:
@staticmethod
def _get_chat_display_name_from_id(chat_id: str) -> str:
"""从chat_id获取显示名称"""
try:
# 首先尝试从chat_stream获取真实群组名称
@@ -795,7 +698,7 @@ class StatisticOutputTask(AsyncTask):
# 移除_generate_versions_tab方法
def _generate_html_report(self, stat: dict[str, Any], now: datetime):
async def _generate_html_report(self, stat: dict[str, Any], now: datetime):
"""
生成HTML格式的统计报告
:param stat: 统计数据
@@ -940,8 +843,8 @@ class StatisticOutputTask(AsyncTask):
)
# 不再添加版本对比内容
# 添加图表内容
chart_data = self._generate_chart_data(stat)
# 添加图表内容 (修正缩进)
chart_data = await self._generate_chart_data(stat)
tab_content_list.append(self._generate_chart_tab(chart_data))
joined_tab_list = "\n".join(tab_list)
@@ -1090,106 +993,90 @@ class StatisticOutputTask(AsyncTask):
with open(self.record_file_path, "w", encoding="utf-8") as f:
f.write(html_template)
def _generate_chart_data(self, stat: dict[str, Any]) -> dict:
"""生成图表数据"""
async def _generate_chart_data(self, stat: dict[str, Any]) -> dict:
"""生成图表数据 (异步)"""
now = datetime.now()
chart_data = {}
chart_data: Dict[str, Any] = {}
# 支持多个时间范围
time_ranges = [
("6h", 6, 10), # 6小时，10分钟间隔
("12h", 12, 15), # 12小时，15分钟间隔
("24h", 24, 15), # 24小时，15分钟间隔
("48h", 48, 30), # 48小时，30分钟间隔
("6h", 6, 10),
("12h", 12, 15),
("24h", 24, 15),
("48h", 48, 30),
]
# 依次处理（数据量不大，避免复杂度；如需可改 gather）
for range_key, hours, interval_minutes in time_ranges:
range_data = self._collect_interval_data(now, hours, interval_minutes)
chart_data[range_key] = range_data
chart_data[range_key] = await self._collect_interval_data(now, hours, interval_minutes)
return chart_data
def _collect_interval_data(self, now: datetime, hours: int, interval_minutes: int) -> dict:
"""收集指定时间范围内每个间隔的数据"""
# 生成时间点
async def _collect_interval_data(self, now: datetime, hours: int, interval_minutes: int) -> dict:
start_time = now - timedelta(hours=hours)
time_points = []
time_points: List[datetime] = []
current_time = start_time
while current_time <= now:
time_points.append(current_time)
current_time += timedelta(minutes=interval_minutes)
# 初始化数据结构
total_cost_data = [0] * len(time_points)
cost_by_model = {}
cost_by_module = {}
message_by_chat = {}
total_cost_data = [0.0] * len(time_points)
cost_by_model: Dict[str, List[float]] = {}
cost_by_module: Dict[str, List[float]] = {}
message_by_chat: Dict[str, List[int]] = {}
time_labels = [t.strftime("%H:%M") for t in time_points]
interval_seconds = interval_minutes * 60
# 查询LLM使用记录
query_start_time = start_time
records = _sync_db_get(
model_class=LLMUsage, filters={"timestamp": {"$gte": query_start_time}}, order_by="-timestamp"
)
for record in records:
# 单次查询 LLMUsage
llm_records = await db_get(
model_class=LLMUsage,
filters={"timestamp": {"$gte": start_time}},
order_by="-timestamp",
) or []
for record in llm_records:
if not isinstance(record, dict) or not record.get("timestamp"):
continue
record_time = record["timestamp"]
# 找到对应的时间间隔索引
if isinstance(record_time, str):
try:
record_time = datetime.fromisoformat(record_time)
except Exception:
continue
time_diff = (record_time - start_time).total_seconds()
interval_index = int(time_diff // interval_seconds)
if 0 <= interval_index < len(time_points):
# 累加总花费数据
idx = int(time_diff // interval_seconds)
if 0 <= idx < len(time_points):
cost = record.get("cost") or 0.0
total_cost_data[interval_index] += cost # type: ignore
# 累加按模型分类的花费
total_cost_data[idx] += cost
model_name = record.get("model_name") or "unknown"
if model_name not in cost_by_model:
cost_by_model[model_name] = [0] * len(time_points)
cost_by_model[model_name][interval_index] += cost
# 累加按模块分类的花费
cost_by_model[model_name] = [0.0] * len(time_points)
cost_by_model[model_name][idx] += cost
request_type = record.get("request_type") or "unknown"
module_name = request_type.split(".")[0] if "." in request_type else request_type
if module_name not in cost_by_module:
cost_by_module[module_name] = [0] * len(time_points)
cost_by_module[module_name][interval_index] += cost
cost_by_module[module_name] = [0.0] * len(time_points)
cost_by_module[module_name][idx] += cost
# 查询消息记录
query_start_timestamp = start_time.timestamp()
records = _sync_db_get(
model_class=Messages, filters={"time": {"$gte": query_start_timestamp}}, order_by="-time"
)
for message in records:
message_time_ts = message["time"]
# 找到对应的时间间隔索引
time_diff = message_time_ts - query_start_timestamp
interval_index = int(time_diff // interval_seconds)
if 0 <= interval_index < len(time_points):
# 确定聊天流名称
chat_name = None
if message.get("chat_info_group_id"):
chat_name = message.get("chat_info_group_name") or f"群{message['chat_info_group_id']}"
elif message.get("user_id"):
chat_name = message.get("user_nickname") or f"用户{message['user_id']}"
# 单次查询 Messages
msg_records = await db_get(
model_class=Messages,
filters={"time": {"$gte": start_time.timestamp()}},
order_by="-time",
) or []
for msg in msg_records:
if not isinstance(msg, dict) or not msg.get("time"):
continue
msg_ts = msg["time"]
time_diff = msg_ts - start_time.timestamp()
idx = int(time_diff // interval_seconds)
if 0 <= idx < len(time_points):
if msg.get("chat_info_group_id"):
chat_name = msg.get("chat_info_group_name") or f"群{msg['chat_info_group_id']}"
elif msg.get("user_id"):
chat_name = msg.get("user_nickname") or f"用户{msg['user_id']}"
else:
continue
if not chat_name:
continue
# 累加消息数
if chat_name not in message_by_chat:
message_by_chat[chat_name] = [0] * len(time_points)
message_by_chat[chat_name][interval_index] += 1
message_by_chat[chat_name][idx] += 1
return {
"time_labels": time_labels,
@@ -1199,7 +1086,8 @@ class StatisticOutputTask(AsyncTask):
"message_by_chat": message_by_chat,
}
def _generate_chart_tab(self, chart_data: dict) -> str:
@staticmethod
def _generate_chart_tab(chart_data: dict) -> str:
# sourcery skip: extract-duplicate-method, move-assign-in-block
"""生成图表选项卡HTML内容"""
@@ -1475,101 +1363,4 @@ class StatisticOutputTask(AsyncTask):
}});
</script>
</div>
"""
class AsyncStatisticOutputTask(AsyncTask):
"""完全异步的统计输出任务 - 更高性能版本"""
def __init__(self, record_file_path: str = "maibot_statistics.html"):
# 延迟0秒启动，运行间隔300秒
super().__init__(task_name="Async Statistics Data Output Task", wait_before_start=0, run_interval=300)
# 直接复用 StatisticOutputTask 的初始化逻辑
temp_stat_task = StatisticOutputTask(record_file_path)
self.name_mapping = temp_stat_task.name_mapping
self.record_file_path = temp_stat_task.record_file_path
self.stat_period = temp_stat_task.stat_period
async def run(self):
"""完全异步执行统计任务"""
async def _async_collect_and_output():
try:
now = datetime.now()
loop = asyncio.get_event_loop()
with concurrent.futures.ThreadPoolExecutor() as executor:
logger.info("正在后台收集统计数据...")
# 数据收集任务
collect_task = asyncio.create_task(
loop.run_in_executor(executor, self._collect_all_statistics, now) # type: ignore
)
stats = await collect_task
logger.info("统计数据收集完成")
# 创建并发的输出任务
output_tasks = [
asyncio.create_task(loop.run_in_executor(executor, self._statistic_console_output, stats, now)), # type: ignore
asyncio.create_task(loop.run_in_executor(executor, self._generate_html_report, stats, now)), # type: ignore
]
# 等待所有输出任务完成
await asyncio.gather(*output_tasks)
logger.info("统计数据后台输出完成")
except Exception as e:
logger.exception(f"后台统计数据输出过程中发生异常：{e}")
# 创建后台任务，立即返回
asyncio.create_task(_async_collect_and_output())
# 复用 StatisticOutputTask 的所有方法
def _collect_all_statistics(self, now: datetime):
return StatisticOutputTask._collect_all_statistics(self, now) # type: ignore
def _statistic_console_output(self, stats: Dict[str, Any], now: datetime):
return StatisticOutputTask._statistic_console_output(self, stats, now) # type: ignore
def _generate_html_report(self, stats: dict[str, Any], now: datetime):
return StatisticOutputTask._generate_html_report(self, stats, now) # type: ignore
# 其他需要的方法也可以类似复用...
@staticmethod
def _collect_model_request_for_period(collect_period: List[Tuple[str, datetime]]) -> Dict[str, Any]:
return StatisticOutputTask._collect_model_request_for_period(collect_period)
@staticmethod
def _collect_online_time_for_period(collect_period: List[Tuple[str, datetime]], now: datetime) -> Dict[str, Any]:
return StatisticOutputTask._collect_online_time_for_period(collect_period, now)
def _collect_message_count_for_period(self, collect_period: List[Tuple[str, datetime]]) -> Dict[str, Any]:
return StatisticOutputTask._collect_message_count_for_period(self, collect_period) # type: ignore
@staticmethod
def _format_total_stat(stats: Dict[str, Any]) -> str:
return StatisticOutputTask._format_total_stat(stats)
@staticmethod
def _format_model_classified_stat(stats: Dict[str, Any]) -> str:
return StatisticOutputTask._format_model_classified_stat(stats)
def _format_chat_stat(self, stats: Dict[str, Any]) -> str:
return StatisticOutputTask._format_chat_stat(self, stats) # type: ignore
def _generate_chart_data(self, stat: dict[str, Any]) -> dict:
return StatisticOutputTask._generate_chart_data(self, stat) # type: ignore
def _collect_interval_data(self, now: datetime, hours: int, interval_minutes: int) -> dict:
return StatisticOutputTask._collect_interval_data(self, now, hours, interval_minutes) # type: ignore
def _generate_chart_tab(self, chart_data: dict) -> str:
return StatisticOutputTask._generate_chart_tab(self, chart_data) # type: ignore
def _get_chat_display_name_from_id(self, chat_id: str) -> str:
return StatisticOutputTask._get_chat_display_name_from_id(self, chat_id) # type: ignore
def _convert_defaultdict_to_dict(self, data):
return StatisticOutputTask._convert_defaultdict_to_dict(self, data) # type: ignore
"""

7
src/chat/utils/utils.py
View File

@@ -7,7 +7,7 @@ import numpy as np
from collections import Counter
from maim_message import UserInfo
from typing import Optional, Tuple, Dict, List, Any
from typing import Optional, Tuple, Dict, List, Any, Coroutine
from src.common.logger import get_logger
from src.common.message_repository import find_messages, count_messages
@@ -540,7 +540,8 @@ def get_western_ratio(paragraph):
return western_count / len(alnum_chars)
def count_messages_between(start_time: float, end_time: float, stream_id: str) -> tuple[int, int]:
def count_messages_between(start_time: float, end_time: float, stream_id: str) -> tuple[int, int] | tuple[
Coroutine[Any, Any, int], int]:
"""计算两个时间点之间的消息数量和文本总长度
Args:
@@ -662,7 +663,7 @@ def get_chat_type_and_target_info(chat_id: str) -> Tuple[bool, Optional[Dict]]:
if person_id:
# get_value is async, so await it directly
person_info_manager = get_person_info_manager()
person_name = person_info_manager.get_value_sync(person_id, "person_name")
person_name = person_info_manager.get_value(person_id, "person_name")
target_info["person_id"] = person_id
target_info["person_name"] = person_name

104
src/chat/utils/utils_image.py
View File

@@ -69,7 +69,7 @@ class ImageManager:
os.makedirs(self.IMAGE_DIR, exist_ok=True)
@staticmethod
def _get_description_from_db(image_hash: str, description_type: str) -> Optional[str]:
async def _get_description_from_db(image_hash: str, description_type: str) -> Optional[str]:
"""从数据库获取图片描述
Args:
@@ -80,22 +80,22 @@ class ImageManager:
Optional[str]: 描述文本，如果不存在则返回None
"""
try:
with get_db_session() as session:
record = session.execute(
async with get_db_session() as session:
record = (await session.execute(
select(ImageDescriptions).where(
and_(
ImageDescriptions.image_description_hash == image_hash,
ImageDescriptions.type == description_type,
)
)
).scalar()
)).scalar()
return record.description if record else None
except Exception as e:
logger.error(f"从数据库获取描述失败 (SQLAlchemy): {str(e)}")
return None
@staticmethod
def _save_description_to_db(image_hash: str, description: str, description_type: str) -> None:
async def _save_description_to_db(image_hash: str, description: str, description_type: str) -> None:
"""保存图片描述到数据库
Args:
@@ -105,16 +105,16 @@ class ImageManager:
"""
try:
current_timestamp = time.time()
with get_db_session() as session:
async with get_db_session() as session:
# 查找现有记录
existing = session.execute(
existing = (await session.execute(
select(ImageDescriptions).where(
and_(
ImageDescriptions.image_description_hash == image_hash,
ImageDescriptions.type == description_type,
)
)
).scalar()
)).scalar()
if existing:
# 更新现有记录
@@ -129,12 +129,13 @@ class ImageManager:
timestamp=current_timestamp,
)
session.add(new_desc)
session.commit()
await session.commit()
# 会在上下文管理器中自动调用
except Exception as e:
logger.error(f"保存描述到数据库失败 (SQLAlchemy): {str(e)}")
async def get_emoji_tag(self, image_base64: str) -> str:
@staticmethod
async def get_emoji_tag(image_base64: str) -> str:
from src.chat.emoji_system.emoji_manager import get_emoji_manager
emoji_manager = get_emoji_manager()
@@ -174,7 +175,7 @@ class ImageManager:
logger.debug(f"查询EmojiManager时出错: {e}")
# 查询ImageDescriptions表的缓存描述
if cached_description := self._get_description_from_db(image_hash, "emoji"):
if cached_description := await self._get_description_from_db(image_hash, "emoji"):
logger.info(f"[缓存命中] 使用ImageDescriptions表中的描述: {cached_description}...")
return f"[表情包：{cached_description}]"
@@ -238,7 +239,7 @@ class ImageManager:
logger.info(f"[emoji识别] 详细描述: {detailed_description}... -> 情感标签: {final_emotion}")
if cached_description := self._get_description_from_db(image_hash, "emoji"):
if cached_description := await self._get_description_from_db(image_hash, "emoji"):
logger.warning(f"虽然生成了描述，但是找到缓存表情包描述: {cached_description}")
return f"[表情包：{cached_description}]"
@@ -260,10 +261,10 @@ class ImageManager:
try:
from src.common.database.sqlalchemy_models import get_db_session
with get_db_session() as session:
existing_img = session.execute(
async with get_db_session() as session:
existing_img = (await session.execute(
select(Images).where(and_(Images.emoji_hash == image_hash, Images.type == "emoji"))
).scalar()
)).scalar()
if existing_img:
existing_img.path = file_path
@@ -278,7 +279,7 @@ class ImageManager:
timestamp=current_timestamp,
)
session.add(new_img)
session.commit()
await session.commit()
except Exception as e:
logger.error(f"保存到Images表失败: {str(e)}")
@@ -288,7 +289,7 @@ class ImageManager:
logger.debug("偷取表情包功能已关闭，跳过保存。")
# 保存最终的情感标签到缓存 (ImageDescriptions表)
self._save_description_to_db(image_hash, final_emotion, "emoji")
await self._save_description_to_db(image_hash, final_emotion, "emoji")
return f"[表情包：{final_emotion}]"
@@ -305,9 +306,9 @@ class ImageManager:
image_bytes = base64.b64decode(image_base64)
image_hash = hashlib.md5(image_bytes).hexdigest()
# 优先检查Images表中是否已有完整的描述
with get_db_session() as session:
existing_image = session.execute(select(Images).where(Images.emoji_hash == image_hash)).scalar()
async with get_db_session() as session:
# 优先检查Images表中是否已有完整的描述
existing_image = (await session.execute(select(Images).where(Images.emoji_hash == image_hash))).scalar()
if existing_image:
# 更新计数
if hasattr(existing_image, "count") and existing_image.count is not None:
@@ -317,34 +318,34 @@ class ImageManager:
# 如果已有描述，直接返回
if existing_image.description:
await session.commit()
logger.debug(f"[缓存命中] 使用Images表中的图片描述: {existing_image.description}...")
return f"[图片：{existing_image.description}]"
if cached_description := self._get_description_from_db(image_hash, "image"):
logger.debug(f"[缓存命中] 使用ImageDescriptions表中的描述: {cached_description}...")
return f"[图片：{cached_description}]"
# 如果没有描述，继续在当前会话中操作
if cached_description := await self._get_description_from_db(image_hash, "image"):
logger.debug(f"[缓存命中] 使用ImageDescriptions表中的描述: {cached_description}...")
return f"[图片：{cached_description}]"
# 调用AI获取描述
image_format = Image.open(io.BytesIO(image_bytes)).format.lower() # type: ignore
prompt = global_config.custom_prompt.image_prompt
logger.info(f"[VLM调用] 为图片生成新描述 (Hash: {image_hash[:8]}...)")
description, _ = await self.vlm.generate_response_for_image(
prompt, image_base64, image_format, temperature=0.4, max_tokens=300
)
# 调用AI获取描述
image_format = Image.open(io.BytesIO(image_bytes)).format.lower() # type: ignore
prompt = global_config.custom_prompt.image_prompt
logger.info(f"[VLM调用] 为图片生成新描述 (Hash: {image_hash[:8]}...)")
description, _ = await self.vlm.generate_response_for_image(
prompt, image_base64, image_format, temperature=0.4, max_tokens=300
)
if description is None:
logger.warning("AI未能生成图片描述")
return "[图片(描述生成失败)]"
if description is None:
logger.warning("AI未能生成图片描述")
return "[图片(描述生成失败)]"
# 保存图片和描述
current_timestamp = time.time()
filename = f"{int(current_timestamp)}_{image_hash[:8]}.{image_format}"
image_dir = os.path.join(self.IMAGE_DIR, "image")
os.makedirs(image_dir, exist_ok=True)
file_path = os.path.join(image_dir, filename)
# 保存图片和描述
current_timestamp = time.time()
filename = f"{int(current_timestamp)}_{image_hash[:8]}.{image_format}"
image_dir = os.path.join(self.IMAGE_DIR, "image")
os.makedirs(image_dir, exist_ok=True)
file_path = os.path.join(image_dir, filename)
try:
# 保存文件
with open(file_path, "wb") as f:
f.write(image_bytes)
@@ -357,7 +358,6 @@ class ImageManager:
existing_image.image_id = str(uuid.uuid4())
if not hasattr(existing_image, "vlm_processed") or existing_image.vlm_processed is None:
existing_image.vlm_processed = True
logger.debug(f"[数据库] 更新已有图片记录: {image_hash[:8]}...")
else:
new_img = Images(
@@ -371,13 +371,15 @@ class ImageManager:
count=1,
)
session.add(new_img)
logger.debug(f"[数据库] 创建新图片记录: {image_hash[:8]}...")
except Exception as e:
logger.error(f"保存图片文件或元数据失败: {str(e)}")
# 保存描述到ImageDescriptions表作为备用缓存
self._save_description_to_db(image_hash, description, "image")
await session.commit()
# 保存描述到ImageDescriptions表作为备用缓存
await self._save_description_to_db(image_hash, description, "image")
logger.info(f"[VLM完成] 图片描述生成: {description}...")
return f"[图片：{description}]"
logger.info(f"[VLM完成] 图片描述生成: {description}...")
return f"[图片：{description}]"
@@ -524,8 +526,8 @@ class ImageManager:
image_base64 = image_base64.encode("ascii", errors="ignore").decode("ascii")
image_bytes = base64.b64decode(image_base64)
image_hash = hashlib.md5(image_bytes).hexdigest()
with get_db_session() as session:
existing_image = session.execute(select(Images).where(Images.emoji_hash == image_hash)).scalar()
async with get_db_session() as session:
existing_image = (await session.execute(select(Images).where(Images.emoji_hash == image_hash))).scalar()
if existing_image:
# 检查是否缺少必要字段，如果缺少则创建新记录
if (
@@ -545,6 +547,7 @@ class ImageManager:
existing_image.vlm_processed = False
existing_image.count += 1
await session.commit()
# 如果已有描述，直接返回
if existing_image.description and existing_image.description.strip():
@@ -555,6 +558,7 @@ class ImageManager:
# 更新数据库中的描述
existing_image.description = description.replace("[图片：", "").replace("]", "")
existing_image.vlm_processed = True
await session.commit()
return existing_image.image_id, f"[picid:{existing_image.image_id}]"
# print(f"图片不存在: {image_hash}")
@@ -587,7 +591,7 @@ class ImageManager:
count=1,
)
session.add(new_img)
session.commit()
await session.commit()
return image_id, f"[picid:{image_id}]"

1089
src/chat/utils/utils_video.py
View File

File diff suppressed because it is too large Load Diff

7
src/chat/utils/utils_video_legacy.py
View File

@@ -13,7 +13,7 @@ import base64
import numpy as np
from PIL import Image
from pathlib import Path
from typing import List, Tuple, Optional
from typing import List, Tuple, Optional, Any
import io
from concurrent.futures import ThreadPoolExecutor
@@ -31,7 +31,7 @@ def _extract_frames_worker(
max_image_size: int,
frame_extraction_mode: str,
frame_interval_seconds: Optional[float],
) -> List[Tuple[str, float]]:
) -> list[Any] | list[tuple[str, str]]:
"""线程池中提取视频帧的工作函数"""
frames = []
try:
@@ -568,7 +568,8 @@ class LegacyVideoAnalyzer:
logger.error(error_msg)
return error_msg
def is_supported_video(self, file_path: str) -> bool:
@staticmethod
def is_supported_video(file_path: str) -> bool:
"""检查是否为支持的视频格式"""
supported_formats = {".mp4", ".avi", ".mov", ".mkv", ".flv", ".wmv", ".m4v", ".3gp", ".webm"}
return Path(file_path).suffix.lower() in supported_formats

6
src/common/cache_manager.py
View File

@@ -53,7 +53,8 @@ class CacheManager:
self._initialized = True
logger.info("缓存管理器已初始化: L1 (内存+FAISS), L2 (数据库+ChromaDB)")
def _validate_embedding(self, embedding_result: Any) -> Optional[np.ndarray]:
@staticmethod
def _validate_embedding(embedding_result: Any) -> Optional[np.ndarray]:
"""
验证和标准化嵌入向量格式
"""
@@ -90,7 +91,8 @@ class CacheManager:
logger.error(f"验证嵌入向量时发生错误: {e}")
return None
def _generate_key(self, tool_name: str, function_args: Dict[str, Any], tool_file_path: Union[str, Path]) -> str:
@staticmethod
def _generate_key(tool_name: str, function_args: Dict[str, Any], tool_file_path: Union[str, Path]) -> str:
"""生成确定性的缓存键，包含文件修改时间以实现自动失效。"""
try:
tool_file_path = Path(tool_file_path)

6
src/common/data_models/info_data_model.py
View File

@@ -1,10 +1,10 @@
from dataclasses import dataclass, field
from typing import Optional, Dict, List, TYPE_CHECKING
from . import BaseDataModel
if TYPE_CHECKING:
from .database_data_model import DatabaseMessages
from src.plugin_system.base.component_types import ActionInfo, ChatMode
pass
@dataclass
@@ -21,7 +21,7 @@ class ActionPlannerInfo(BaseDataModel):
action_type: str = field(default_factory=str)
reasoning: Optional[str] = None
action_data: Optional[Dict] = None
action_message: Optional["DatabaseMessages"] = None
action_message: Optional[Dict] = None
available_actions: Optional[Dict[str, "ActionInfo"]] = None

3
src/common/data_models/llm_data_model.py
View File

@@ -2,8 +2,9 @@ from dataclasses import dataclass
from typing import Optional, List, Tuple, TYPE_CHECKING, Any
from . import BaseDataModel
if TYPE_CHECKING:
from src.llm_models.payload_content.tool_option import ToolCall
pass
@dataclass
class LLMGenerationDataModel(BaseDataModel):

4
src/common/data_models/message_data_model.py
View File

@@ -1,10 +1,10 @@
from typing import Optional, TYPE_CHECKING
from dataclasses import dataclass, field
from typing import Optional, TYPE_CHECKING
from . import BaseDataModel
if TYPE_CHECKING:
from .database_data_model import DatabaseMessages
pass
@dataclass

36
src/common/database/database.py
View File

@@ -25,27 +25,39 @@ class DatabaseProxy:
self._engine = None
self._session = None
def initialize(self, *args, **kwargs):
@staticmethod
def initialize(*args, **kwargs):
"""初始化数据库连接"""
return initialize_database_compat()
class SQLAlchemyTransaction:
"""SQLAlchemy事务上下文管理器"""
"""SQLAlchemy 异步事务上下文管理器 (兼容旧代码示例，推荐直接使用 get_db_session)。"""
def __init__(self):
self._ctx = None
self.session = None
def __enter__(self):
self.session = get_db_session()
async def __aenter__(self):
# get_db_session 是一个 async contextmanager
self._ctx = get_db_session()
self.session = await self._ctx.__aenter__()
return self.session
def __exit__(self, exc_type, exc_val, exc_tb):
if exc_type is None:
self.session.commit()
else:
self.session.rollback()
self.session.close()
async def __aexit__(self, exc_type, exc_val, exc_tb):
try:
if self.session:
if exc_type is None:
try:
await self.session.commit()
except Exception:
await self.session.rollback()
raise
else:
await self.session.rollback()
finally:
if self._ctx:
await self._ctx.__aexit__(exc_type, exc_val, exc_tb)
# 创建全局数据库代理实例
@@ -89,7 +101,7 @@ def get_db():
return _db
def initialize_sql_database(database_config):
async def initialize_sql_database(database_config):
"""
根据配置初始化SQL数据库连接（SQLAlchemy版本）
@@ -119,7 +131,7 @@ def initialize_sql_database(database_config):
# 使用SQLAlchemy初始化
success = initialize_database_compat()
if success:
_sql_engine = get_engine()
_sql_engine = await get_engine()
logger.info("SQLAlchemy数据库初始化成功")
else:
logger.error("SQLAlchemy数据库初始化失败")

139
src/common/database/db_migration.py
View File

@@ -1,77 +1,116 @@
# mmc/src/common/database/db_migration.py
from sqlalchemy import inspect, text
from sqlalchemy import inspect
from sqlalchemy.schema import AddColumn, CreateIndex
from src.common.database.sqlalchemy_models import Base, get_engine
from src.common.logger import get_logger
logger = get_logger("db_migration")
def check_and_migrate_database():
async def check_and_migrate_database():
"""
检查数据库结构并自动迁移（添加缺失的表和列）。
异步检查数据库结构并自动迁移。
- 自动创建不存在的表。
- 自动为现有表添加缺失的列。
- 自动为现有表创建缺失的索引。
"""
logger.info("正在检查数据库结构并执行自动迁移...")
engine = get_engine()
inspector = inspect(engine)
engine = await get_engine()
# 1. 获取数据库中所有已存在的表名
db_table_names = set(inspector.get_table_names())
async with engine.connect() as connection:
# 在同步上下文中运行inspector操作
def get_inspector(sync_conn):
return inspect(sync_conn)
# 2. 遍历所有在代码中定义的模型
for table_name, table in Base.metadata.tables.items():
logger.debug(f"正在检查表: {table_name}")
inspector = await connection.run_sync(get_inspector)
# 3. 如果表不存在，则创建它
if table_name not in db_table_names:
logger.info(f"表 '{table_name}' 不存在，正在创建...")
# 在同步lambda中传递inspector
db_table_names = await connection.run_sync(lambda conn: set(inspector.get_table_names(conn)))
# 1. 首先处理表的创建
tables_to_create = []
for table_name, table in Base.metadata.tables.items():
if table_name not in db_table_names:
tables_to_create.append(table)
if tables_to_create:
logger.info(f"发现 {len(tables_to_create)} 个不存在的表，正在创建...")
try:
table.create(engine)
logger.info(f"表 '{table_name}' 创建成功。")
# 一次性创建所有缺失的表
await connection.run_sync(
lambda sync_conn: Base.metadata.create_all(sync_conn, tables=tables_to_create)
)
for table in tables_to_create:
logger.info(f"表 '{table.name}' 创建成功。")
db_table_names.add(table.name) # 将新创建的表添加到集合中
except Exception as e:
logger.error(f"创建表 '{table_name}' 失败: {e}")
continue
logger.error(f"创建表时失败: {e}", exc_info=True)
# 4. 如果表已存在，则检查并添加缺失的列
db_columns = {col["name"] for col in inspector.get_columns(table_name)}
model_columns = {col.name for col in table.c}
# 2. 然后处理现有表的列和索引的添加
for table_name, table in Base.metadata.tables.items():
if table_name not in db_table_names:
logger.warning(f"跳过检查表 '{table_name}'，因为它在创建步骤中可能已失败。")
continue
missing_columns = model_columns - db_columns
if not missing_columns:
logger.debug(f"表 '{table_name}' 结构一致，无需修改。")
continue
logger.debug(f"正在检查表 '{table_name}' 的列和索引...")
logger.info(f"在表 '{table_name}' 中发现缺失的列: {', '.join(missing_columns)}")
with engine.connect() as connection:
trans = connection.begin()
try:
for column_name in missing_columns:
column = table.c[column_name]
# 检查并添加缺失的列
db_columns = await connection.run_sync(
lambda conn: {col["name"] for col in inspector.get_columns(table_name, conn)}
)
model_columns = {col.name for col in table.c}
missing_columns = model_columns - db_columns
# 构造并执行 ALTER TABLE 语句
try:
column_type = column.type.compile(engine.dialect)
sql = f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_type}"
if missing_columns:
logger.info(f"在表 '{table_name}' 中发现缺失的列: {', '.join(missing_columns)}")
async with connection.begin() as trans:
for column_name in missing_columns:
try:
column = table.c[column_name]
add_column_ddl = AddColumn(table_name, column)
await connection.execute(add_column_ddl)
logger.info(f"成功向表 '{table_name}' 添加列 '{column_name}'。")
except Exception as e:
logger.error(
f"向表 '{table_name}' 添加列 '{column_name}' 失败: {e}",
exc_info=True,
)
await trans.rollback()
break # 如果一列失败，则停止处理此表的其他列
else:
logger.info(f"表 '{table_name}' 的列结构一致。")
# 添加默认值和非空约束的处理
if column.default is not None:
default_value = column.default.arg
if isinstance(default_value, str):
sql += f" DEFAULT '{default_value}'"
else:
sql += f" DEFAULT {default_value}"
# 检查并创建缺失的索引
db_indexes = await connection.run_sync(
lambda conn: {idx["name"] for idx in inspector.get_indexes(table_name, conn)}
)
model_indexes = {idx.name for idx in table.indexes}
missing_indexes = model_indexes - db_indexes
if not column.nullable:
sql += " NOT NULL"
if missing_indexes:
logger.info(f"在表 '{table_name}' 中发现缺失的索引: {', '.join(missing_indexes)}")
async with connection.begin() as trans:
for index_name in missing_indexes:
try:
index_obj = next((idx for idx in table.indexes if idx.name == index_name), None)
if index_obj is not None:
await connection.execute(CreateIndex(index_obj))
logger.info(f"成功为表 '{table_name}' 创建索引 '{index_name}'。")
except Exception as e:
logger.error(
f"为表 '{table_name}' 创建索引 '{index_name}' 失败: {e}",
exc_info=True,
)
await trans.rollback()
break # 如果一个索引失败，则停止处理此表的其他索引
else:
logger.debug(f"表 '{table_name}' 的索引一致。")
connection.execute(text(sql))
logger.info(f"成功向表 '{table_name}' 添加列 '{column_name}'。")
except Exception as e:
logger.error(f"向表 '{table_name}' 添加列 '{column_name}' 失败: {e}")
trans.commit()
except Exception as e:
logger.error(f"在表 '{table_name}' 添加列时发生错误，事务已回滚: {e}")
trans.rollback()
logger.error(f"在处理表 '{table_name}' 时发生意外错误: {e}", exc_info=True)
continue
logger.info("数据库结构检查与自动迁移完成。")

108
src/common/database/sqlalchemy_database_api.py
View File

@@ -4,14 +4,14 @@
支持自动重连、连接池管理和更好的错误处理
"""
import traceback
import time
from typing import Dict, List, Any, Union, Type, Optional
import traceback
from typing import Dict, List, Any, Union, Optional
from sqlalchemy import desc, asc, func, and_, select
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy import desc, asc, func, and_
from src.common.logger import get_logger
from src.common.database.sqlalchemy_models import (
Base,
get_db_session,
Messages,
ActionRecords,
@@ -31,6 +31,7 @@ from src.common.database.sqlalchemy_models import (
MaiZoneScheduleStatus,
CacheEntries,
)
from src.common.logger import get_logger
logger = get_logger("sqlalchemy_database_api")
@@ -56,7 +57,7 @@ MODEL_MAPPING = {
}
def build_filters(session, model_class: Type[Base], filters: Dict[str, Any]):
async def build_filters(model_class, filters: Dict[str, Any]):
"""构建查询过滤条件"""
conditions = []
@@ -94,7 +95,7 @@ def build_filters(session, model_class: Type[Base], filters: Dict[str, Any]):
async def db_query(
model_class: Type[Base],
model_class,
data: Optional[Dict[str, Any]] = None,
query_type: Optional[str] = "get",
filters: Optional[Dict[str, Any]] = None,
@@ -102,7 +103,7 @@ async def db_query(
order_by: Optional[List[str]] = None,
single_result: Optional[bool] = False,
) -> Union[List[Dict[str, Any]], Dict[str, Any], None]:
"""执行数据库查询操作
"""执行异步数据库查询操作
Args:
model_class: SQLAlchemy模型类
@@ -120,15 +121,15 @@ async def db_query(
if query_type not in ["get", "create", "update", "delete", "count"]:
raise ValueError("query_type must be 'get', 'create', 'update', 'delete' or 'count'")
with get_db_session() as session:
async with get_db_session() as session:
if query_type == "get":
query = session.query(model_class)
query = select(model_class)
# 应用过滤条件
if filters:
conditions = build_filters(session, model_class, filters)
conditions = await build_filters(model_class, filters)
if conditions:
query = query.filter(and_(*conditions))
query = query.where(and_(*conditions))
# 应用排序
if order_by:
@@ -146,14 +147,15 @@ async def db_query(
query = query.limit(limit)
# 执行查询
results = query.all()
result = await session.execute(query)
results = result.scalars().all()
# 转换为字典格式
result_dicts = []
for result in results:
for result_obj in results:
result_dict = {}
for column in result.__table__.columns:
result_dict[column.name] = getattr(result, column.name)
for column in result_obj.__table__.columns:
result_dict[column.name] = getattr(result_obj, column.name)
result_dicts.append(result_dict)
if single_result:
@@ -167,7 +169,7 @@ async def db_query(
# 创建新记录
new_record = model_class(**data)
session.add(new_record)
session.flush() # 获取自动生成的ID
await session.flush() # 获取自动生成的ID
# 转换为字典格式返回
result_dict = {}
@@ -179,43 +181,60 @@ async def db_query(
if not data:
raise ValueError("更新记录需要提供data参数")
query = session.query(model_class)
query = select(model_class)
# 应用过滤条件
if filters:
conditions = build_filters(session, model_class, filters)
conditions = await build_filters(model_class, filters)
if conditions:
query = query.filter(and_(*conditions))
query = query.where(and_(*conditions))
# 执行更新
affected_rows = query.update(data)
# 首先获取要更新的记录
result = await session.execute(query)
records_to_update = result.scalars().all()
# 更新每个记录
affected_rows = 0
for record in records_to_update:
for field, value in data.items():
if hasattr(record, field):
setattr(record, field, value)
affected_rows += 1
return affected_rows
elif query_type == "delete":
query = session.query(model_class)
query = select(model_class)
# 应用过滤条件
if filters:
conditions = build_filters(session, model_class, filters)
conditions = await build_filters(model_class, filters)
if conditions:
query = query.filter(and_(*conditions))
query = query.where(and_(*conditions))
# 执行删除
affected_rows = query.delete()
# 首先获取要删除的记录
result = await session.execute(query)
records_to_delete = result.scalars().all()
# 删除记录
affected_rows = 0
for record in records_to_delete:
session.delete(record)
affected_rows += 1
return affected_rows
elif query_type == "count":
query = session.query(func.count(model_class.id))
query = select(func.count(model_class.id))
# 应用过滤条件
if filters:
base_query = session.query(model_class)
conditions = build_filters(session, model_class, filters)
conditions = await build_filters(model_class, filters)
if conditions:
base_query = base_query.filter(and_(*conditions))
query = session.query(func.count()).select_from(base_query.subquery())
query = query.where(and_(*conditions))
return query.scalar()
result = await session.execute(query)
return result.scalar()
except SQLAlchemyError as e:
logger.error(f"[SQLAlchemy] 数据库操作出错: {e}")
@@ -238,9 +257,9 @@ async def db_query(
async def db_save(
model_class: Type[Base], data: Dict[str, Any], key_field: Optional[str] = None, key_value: Optional[Any] = None
model_class, data: Dict[str, Any], key_field: Optional[str] = None, key_value: Optional[Any] = None
) -> Optional[Dict[str, Any]]:
"""保存数据到数据库（创建或更新）
"""异步保存数据到数据库（创建或更新）
Args:
model_class: SQLAlchemy模型类
@@ -252,13 +271,13 @@ async def db_save(
保存后的记录数据或None
"""
try:
with get_db_session() as session:
async with get_db_session() as session:
# 如果提供了key_field和key_value，尝试更新现有记录
if key_field and key_value is not None:
if hasattr(model_class, key_field):
existing_record = (
session.query(model_class).filter(getattr(model_class, key_field) == key_value).first()
)
query = select(model_class).where(getattr(model_class, key_field) == key_value)
result = await session.execute(query)
existing_record = result.scalars().first()
if existing_record:
# 更新现有记录
@@ -266,7 +285,7 @@ async def db_save(
if hasattr(existing_record, field):
setattr(existing_record, field, value)
session.flush()
await session.flush()
# 转换为字典格式返回
result_dict = {}
@@ -277,8 +296,7 @@ async def db_save(
# 创建新记录
new_record = model_class(**data)
session.add(new_record)
session.commit()
session.flush()
await session.flush()
# 转换为字典格式返回
result_dict = {}
@@ -297,13 +315,13 @@ async def db_save(
async def db_get(
model_class: Type[Base],
model_class,
filters: Optional[Dict[str, Any]] = None,
limit: Optional[int] = None,
order_by: Optional[str] = None,
single_result: Optional[bool] = False,
) -> Union[List[Dict[str, Any]], Dict[str, Any], None]:
"""从数据库获取记录
"""异步从数据库获取记录
Args:
model_class: SQLAlchemy模型类
@@ -335,7 +353,7 @@ async def store_action_info(
action_data: Optional[dict] = None,
action_name: str = "",
) -> Optional[Dict[str, Any]]:
"""存储动作信息到数据库
"""异步存储动作信息到数据库
Args:
chat_stream: 聊天流对象

37
src/common/database/sqlalchemy_init.py
View File

@@ -1,7 +1,7 @@
"""SQLAlchemy数据库初始化模块
替换Peewee的数据库初始化逻辑
提供统一的数据库初始化接口
提供统一的异步数据库初始化接口
"""
from typing import Optional
@@ -12,25 +12,25 @@ from src.common.database.sqlalchemy_models import Base, get_engine, initialize_d
logger = get_logger("sqlalchemy_init")
def initialize_sqlalchemy_database() -> bool:
async def initialize_sqlalchemy_database() -> bool:
"""
初始化SQLAlchemy数据库
初始化SQLAlchemy异步数据库
创建所有表结构
Returns:
bool: 初始化是否成功
"""
try:
logger.info("开始初始化SQLAlchemy数据库...")
logger.info("开始初始化SQLAlchemy异步数据库...")
# 初始化数据库引擎和会话
engine, session_local = initialize_database()
engine, session_local = await initialize_database()
if engine is None:
logger.error("数据库引擎初始化失败")
return False
logger.info("SQLAlchemy数据库初始化成功")
logger.info("SQLAlchemy异步数据库初始化成功")
return True
except SQLAlchemyError as e:
@@ -41,9 +41,9 @@ def initialize_sqlalchemy_database() -> bool:
return False
def create_all_tables() -> bool:
async def create_all_tables() -> bool:
"""
创建所有数据库表
异步创建所有数据库表
Returns:
bool: 创建是否成功
@@ -51,13 +51,14 @@ def create_all_tables() -> bool:
try:
logger.info("开始创建数据库表...")
engine = get_engine()
engine = await get_engine()
if engine is None:
logger.error("无法获取数据库引擎")
return False
# 创建所有表
Base.metadata.create_all(bind=engine)
# 异步创建所有表
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
logger.info("数据库表创建成功")
return True
@@ -70,15 +71,15 @@ def create_all_tables() -> bool:
return False
def get_database_info() -> Optional[dict]:
async def get_database_info() -> Optional[dict]:
"""
获取数据库信息
异步获取数据库信息
Returns:
dict: 数据库信息字典，包含引擎信息等
"""
try:
engine = get_engine()
engine = await get_engine()
if engine is None:
return None
@@ -100,9 +101,9 @@ def get_database_info() -> Optional[dict]:
_database_initialized = False
def initialize_database_compat() -> bool:
async def initialize_database_compat() -> bool:
"""
兼容性数据库初始化函数
兼容性异步数据库初始化函数
用于替换原有的Peewee初始化代码
Returns:
@@ -113,9 +114,9 @@ def initialize_database_compat() -> bool:
if _database_initialized:
return True
success = initialize_sqlalchemy_database()
success = await initialize_sqlalchemy_database()
if success:
success = create_all_tables()
success = await create_all_tables()
if success:
_database_initialized = True

61
src/common/database/sqlalchemy_models.py
View File

@@ -3,16 +3,18 @@
替换Peewee ORM，使用SQLAlchemy提供更好的连接池管理和错误恢复能力
"""
from sqlalchemy import Column, String, Float, Integer, Boolean, Text, Index, create_engine, DateTime
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, Session, Mapped, mapped_column
from sqlalchemy.pool import QueuePool
import os
import datetime
import os
import time
from typing import Iterator, Optional, Any, Dict
from contextlib import asynccontextmanager
from typing import Optional, Any, Dict, AsyncGenerator
from sqlalchemy import Column, String, Float, Integer, Boolean, Text, Index, DateTime
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import Mapped, mapped_column
from src.common.logger import get_logger
from contextlib import contextmanager
logger = get_logger("sqlalchemy_models")
@@ -575,14 +577,14 @@ def get_database_url():
# 使用Unix socket连接
encoded_socket = quote_plus(config.mysql_unix_socket)
return (
f"mysql+pymysql://{encoded_user}:{encoded_password}"
f"mysql+aiomysql://{encoded_user}:{encoded_password}"
f"@/{config.mysql_database}"
f"?unix_socket={encoded_socket}&charset={config.mysql_charset}"
)
else:
# 使用标准TCP连接
return (
f"mysql+pymysql://{encoded_user}:{encoded_password}"
f"mysql+aiomysql://{encoded_user}:{encoded_password}"
f"@{config.mysql_host}:{config.mysql_port}/{config.mysql_database}"
f"?charset={config.mysql_charset}"
)
@@ -597,11 +599,11 @@ def get_database_url():
# 确保数据库目录存在
os.makedirs(os.path.dirname(db_path), exist_ok=True)
return f"sqlite:///{db_path}"
return f"sqlite+aiosqlite:///{db_path}"
def initialize_database():
"""初始化数据库引擎和会话"""
async def initialize_database():
"""初始化异步数据库引擎和会话"""
global _engine, _SessionLocal
if _engine is not None:
@@ -619,10 +621,9 @@ def initialize_database():
}
if config.database_type == "mysql":
# MySQL连接池配置
# MySQL连接池配置 - 异步引擎使用默认连接池
engine_kwargs.update(
{
"poolclass": QueuePool,
"pool_size": config.connection_pool_size,
"max_overflow": config.connection_pool_size * 2,
"pool_timeout": config.connection_timeout,
@@ -638,10 +639,9 @@ def initialize_database():
}
)
else:
# SQLite配置 - 添加连接池设置以避免连接耗尽
# SQLite配置 - 异步引擎使用默认连接池
engine_kwargs.update(
{
"poolclass": QueuePool,
"pool_size": 20, # 增加池大小
"max_overflow": 30, # 增加溢出连接数
"pool_timeout": 60, # 增加超时时间
@@ -654,41 +654,40 @@ def initialize_database():
}
)
_engine = create_engine(database_url, **engine_kwargs)
_SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=_engine)
_engine = create_async_engine(database_url, **engine_kwargs)
_SessionLocal = async_sessionmaker(bind=_engine, class_=AsyncSession, expire_on_commit=False)
# 调用新的迁移函数，它会处理表的创建和列的添加
from src.common.database.db_migration import check_and_migrate_database
check_and_migrate_database()
await check_and_migrate_database()
logger.info(f"SQLAlchemy数据库初始化成功: {config.database_type}")
logger.info(f"SQLAlchemy异步数据库初始化成功: {config.database_type}")
return _engine, _SessionLocal
@contextmanager
def get_db_session() -> Iterator[Session]:
"""数据库会话上下文管理器 - 推荐使用这个而不是get_session()"""
session: Optional[Session] = None
@asynccontextmanager
async def get_db_session() -> AsyncGenerator[AsyncSession, None]:
"""异步数据库会话上下文管理器"""
session: Optional[AsyncSession] = None
try:
engine, SessionLocal = initialize_database()
engine, SessionLocal = await initialize_database()
if not SessionLocal:
raise RuntimeError("Database session not initialized")
session = SessionLocal()
yield session
# session.commit()
except Exception:
if session:
session.rollback()
await session.rollback()
raise
finally:
if session:
session.close()
await session.close()
def get_engine():
"""获取数据库引擎"""
engine, _ = initialize_database()
async def get_engine():
"""获取异步数据库引擎"""
engine, _ = await initialize_database()
return engine

20
src/common/message_repository.py
View File

@@ -25,7 +25,7 @@ def _model_to_dict(instance: Base) -> Dict[str, Any]:
return {col.name: getattr(instance, col.name) for col in instance.__table__.columns}
def find_messages(
async def find_messages(
message_filter: dict[str, Any],
sort: Optional[List[tuple[str, int]]] = None,
limit: int = 0,
@@ -46,7 +46,7 @@ def find_messages(
消息字典列表，如果出错则返回空列表。
"""
try:
with get_db_session() as session:
async with get_db_session() as session:
query = select(Messages)
# 应用过滤器
@@ -96,7 +96,7 @@ def find_messages(
# 获取时间最早的 limit 条记录，已经是正序
query = query.order_by(Messages.time.asc()).limit(limit)
try:
results = session.execute(query).scalars().all()
results = (await session.execute(query)).scalars().all()
except Exception as e:
logger.error(f"执行earliest查询失败: {e}")
results = []
@@ -104,7 +104,7 @@ def find_messages(
# 获取时间最晚的 limit 条记录
query = query.order_by(Messages.time.desc()).limit(limit)
try:
latest_results = session.execute(query).scalars().all()
latest_results = (await session.execute(query)).scalars().all()
# 将结果按时间正序排列
results = sorted(latest_results, key=lambda msg: msg.time)
except Exception as e:
@@ -128,12 +128,12 @@ def find_messages(
if sort_terms:
query = query.order_by(*sort_terms)
try:
results = session.execute(query).scalars().all()
results = (await session.execute(query)).scalars().all()
except Exception as e:
logger.error(f"执行无限制查询失败: {e}")
results = []
return [_model_to_dict(msg) for msg in results]
return [_model_to_dict(msg) for msg in results]
except Exception as e:
log_message = (
f"使用 SQLAlchemy 查找消息失败 (filter={message_filter}, sort={sort}, limit={limit}, limit_mode={limit_mode}): {e}\n"
@@ -143,7 +143,7 @@ def find_messages(
return []
def count_messages(message_filter: dict[str, Any]) -> int:
async def count_messages(message_filter: dict[str, Any]) -> int:
"""
根据提供的过滤器计算消息数量。
@@ -154,7 +154,7 @@ def count_messages(message_filter: dict[str, Any]) -> int:
符合条件的消息数量，如果出错则返回 0。
"""
try:
with get_db_session() as session:
async with get_db_session() as session:
query = select(func.count(Messages.id))
# 应用过滤器
@@ -192,7 +192,7 @@ def count_messages(message_filter: dict[str, Any]) -> int:
if conditions:
query = query.where(*conditions)
count = session.execute(query).scalar()
count = (await session.execute(query)).scalar()
return count or 0
except Exception as e:
log_message = f"使用 SQLAlchemy 计数消息失败 (message_filter={message_filter}): {e}\n{traceback.format_exc()}"
@@ -201,5 +201,5 @@ def count_messages(message_filter: dict[str, Any]) -> int:
# 你可以在这里添加更多与 messages 集合相关的数据库操作函数，例如 find_one_message, insert_message 等。
# 注意：对于 SQLAlchemy，插入操作通常是使用 session.add() 和 session.commit()。
# 注意：对于 SQLAlchemy，插入操作通常是使用 session.add() 和 await session.commit()。
# 查找单个消息可以使用 session.execute(select(Messages).where(...)).scalar_one_or_none()。

8
src/common/server.py
View File

@@ -1,10 +1,10 @@
import os
from typing import Optional
from fastapi import FastAPI, APIRouter
from fastapi.middleware.cors import CORSMiddleware # 新增导入
from typing import Optional
from uvicorn import Config, Server as UvicornServer
from src.config.config import global_config
from rich.traceback import install
import os
from uvicorn import Config, Server as UvicornServer
install(extra_lines=3)

8
src/config/api_ada_configs.py
View File

@@ -22,7 +22,6 @@ class APIProvider(ValidatedConfigBase):
enable_content_obfuscation: bool = Field(default=False, description="是否启用内容混淆（用于特定场景下的内容处理）")
obfuscation_intensity: int = Field(default=1, ge=1, le=3, description="混淆强度（1-3级，数值越高混淆程度越强）")
@field_validator("base_url")
@classmethod
def validate_base_url(cls, v):
"""验证base_url，确保URL格式正确"""
@@ -30,7 +29,6 @@ class APIProvider(ValidatedConfigBase):
raise ValueError("base_url必须以http://或https://开头")
return v
@field_validator("api_key")
@classmethod
def validate_api_key(cls, v):
"""验证API密钥不能为空"""
@@ -75,7 +73,6 @@ class ModelInfo(ValidatedConfigBase):
extra_params: Dict[str, Any] = Field(default_factory=dict, description="额外参数（用于API调用时的额外配置）")
anti_truncation: bool = Field(default=False, description="是否启用反截断功能，防止模型输出被截断")
@field_validator("price_in", "price_out")
@classmethod
def validate_prices(cls, v):
"""验证价格必须为非负数"""
@@ -83,7 +80,6 @@ class ModelInfo(ValidatedConfigBase):
raise ValueError("价格不能为负数")
return v
@field_validator("model_identifier")
@classmethod
def validate_model_identifier(cls, v):
"""验证模型标识符不能为空且不能包含特殊字符"""
@@ -94,7 +90,6 @@ class ModelInfo(ValidatedConfigBase):
raise ValueError("模型标识符不能包含空格或换行符")
return v
@field_validator("name")
@classmethod
def validate_name(cls, v):
"""验证模型名称不能为空"""
@@ -111,7 +106,6 @@ class TaskConfig(ValidatedConfigBase):
temperature: float = Field(default=0.7, description="模型温度")
concurrency_count: int = Field(default=1, description="并发请求数量")
@field_validator("model_list")
@classmethod
def validate_model_list(cls, v):
"""验证模型列表不能为空"""
@@ -178,7 +172,6 @@ class APIAdapterConfig(ValidatedConfigBase):
self.api_providers_dict = {provider.name: provider for provider in self.api_providers}
self.models_dict = {model.name: model for model in self.models}
@field_validator("models")
@classmethod
def validate_models_list(cls, v):
"""验证模型列表"""
@@ -197,7 +190,6 @@ class APIAdapterConfig(ValidatedConfigBase):
return v
@field_validator("api_providers")
@classmethod
def validate_api_providers_list(cls, v):
"""验证API提供商列表"""

2
src/config/config.py
View File

@@ -412,7 +412,6 @@ class APIAdapterConfig(ValidatedConfigBase):
self.api_providers_dict = {provider.name: provider for provider in self.api_providers}
self.models_dict = {model.name: model for model in self.models}
@field_validator("models")
@classmethod
def validate_models_list(cls, v):
"""验证模型列表"""
@@ -431,7 +430,6 @@ class APIAdapterConfig(ValidatedConfigBase):
return v
@field_validator("api_providers")
@classmethod
def validate_api_providers_list(cls, v):
"""验证API提供商列表"""

2
src/config/config_base.py
View File

@@ -50,7 +50,7 @@ class ConfigBase:
except Exception as e:
raise RuntimeError(f"Failed to convert field '{field_name}' to target type: {e}") from e
return cls(**init_args)
return cls()
@classmethod
def _convert_field(cls, value: Any, field_type: Type[Any]) -> Any:

39
src/config/official_configs.py
View File

@@ -57,6 +57,36 @@ class PersonalityConfig(ValidatedConfigBase):
prompt_mode: Literal["s4u", "normal"] = Field(default="s4u", description="Prompt模式")
compress_personality: bool = Field(default=True, description="是否压缩人格")
compress_identity: bool = Field(default=True, description="是否压缩身份")
# 回复规则配置
reply_targeting_rules: List[str] = Field(
default_factory=lambda: [
"拒绝任何包含骚扰、冒犯、暴力、色情或危险内容的请求。",
"在拒绝时，请使用符合你人设的、坚定的语气。",
"不要执行任何可能被用于恶意目的的指令。"
],
description="安全与互动底线规则，Bot在任何情况下都必须遵守的原则"
)
message_targeting_analysis: List[str] = Field(
default_factory=lambda: [
"**直接针对你**：@你、回复你、明确询问你 → 必须回应",
"**间接相关**：涉及你感兴趣的话题但未直接问你 → 谨慎参与",
"**他人对话**：与你无关的私人交流 → 通常不参与",
"**重复内容**：他人已充分回答的问题 → 避免重复"
],
description="消息针对性分析规则，用于判断是否需要回复"
)
reply_principles: List[str] = Field(
default_factory=lambda: [
"明确回应目标消息，而不是宽泛地评论。",
"可以分享你的看法、提出相关问题，或者开个合适的玩笑。",
"目的是让对话更有趣、更深入。",
"不要浮夸，不要夸张修辞，不要输出多余内容(包括前后缀，冒号和引号，括号()，表情包，at或 @等 )。"
],
description="回复原则，指导如何回复消息"
)
class RelationshipConfig(ValidatedConfigBase):
@@ -122,7 +152,8 @@ class ChatConfig(ValidatedConfigBase):
global_frequency = self._get_global_frequency()
return self.talk_frequency if global_frequency is None else global_frequency
def _get_time_based_frequency(self, time_freq_list: list[str]) -> Optional[float]:
@staticmethod
def _get_time_based_frequency(time_freq_list: list[str]) -> Optional[float]:
"""
根据时间配置列表获取当前时段的频率
@@ -201,7 +232,8 @@ class ChatConfig(ValidatedConfigBase):
return None
def _parse_stream_config_to_chat_id(self, stream_config_str: str) -> Optional[str]:
@staticmethod
def _parse_stream_config_to_chat_id(stream_config_str: str) -> Optional[str]:
"""
解析流配置字符串并生成对应的 chat_id
@@ -280,7 +312,8 @@ class ExpressionConfig(ValidatedConfigBase):
rules: List[ExpressionRule] = Field(default_factory=list, description="表达学习规则")
def _parse_stream_config_to_chat_id(self, stream_config_str: str) -> Optional[str]:
@staticmethod
def _parse_stream_config_to_chat_id(stream_config_str: str) -> Optional[str]:
"""
解析流配置字符串并生成对应的 chat_id

3
src/individuality/individuality.py
View File

@@ -94,8 +94,9 @@ class Individuality:
prompt_personality = f"{personality}\n{identity}"
return f"你的名字是{bot_name}{bot_nickname}，你{prompt_personality}"
@staticmethod
def _get_config_hash(
self, bot_nickname: str, personality_core: str, personality_side: str, identity: str
bot_nickname: str, personality_core: str, personality_side: str, identity: str
) -> tuple[str, str]:
"""获取personality和identity配置的哈希值

4
src/llm_models/payload_content/message.py
View File

@@ -58,7 +58,7 @@ class MessageBuilder:
self,
image_format: str,
image_base64: str,
support_formats: list[str] = SUPPORTED_IMAGE_FORMATS, # 默认支持格式
support_formats=None, # 默认支持格式
) -> "MessageBuilder":
"""
添加图片内容
@@ -66,6 +66,8 @@ class MessageBuilder:
:param image_base64: 图片的base64编码
:return: MessageBuilder对象
"""
if support_formats is None:
support_formats = SUPPORTED_IMAGE_FORMATS
if image_format.lower() not in support_formats:
raise ValueError("不受支持的图片格式")
if not image_base64:

10
src/llm_models/utils.py
View File

@@ -145,9 +145,9 @@ class LLMUsageRecorder:
LLM使用情况记录器（SQLAlchemy版本）
"""
def record_usage_to_database(
self,
model_info: ModelInfo,
@staticmethod
async def record_usage_to_database(
model_info: ModelInfo,
model_usage: UsageRecord,
user_id: str,
request_type: str,
@@ -161,7 +161,7 @@ class LLMUsageRecorder:
session = None
try:
# 使用 SQLAlchemy 会话创建记录
with get_db_session() as session:
async with get_db_session() as session:
usage_record = LLMUsage(
model_name=model_info.model_identifier,
model_assign_name=model_info.name,
@@ -179,7 +179,7 @@ class LLMUsageRecorder:
)
session.add(usage_record)
session.commit()
await session.commit()
logger.debug(
f"Token使用情况 - 模型: {model_usage.model_name}, "

16
src/llm_models/utils_model.py
View File

@@ -202,7 +202,7 @@ class LLMRequest:
content, extracted_reasoning = self._extract_reasoning(content)
reasoning_content = extracted_reasoning
if usage := response.usage:
llm_usage_recorder.record_usage_to_database(
await llm_usage_recorder.record_usage_to_database(
model_info=model_info,
model_usage=usage,
user_id="system",
@@ -367,7 +367,7 @@ class LLMRequest:
# 成功获取响应
if usage := response.usage:
llm_usage_recorder.record_usage_to_database(
await llm_usage_recorder.record_usage_to_database(
model_info=model_info,
model_usage=usage,
time_cost=time.time() - start_time,
@@ -442,7 +442,7 @@ class LLMRequest:
embedding = response.embedding
if usage := response.usage:
llm_usage_recorder.record_usage_to_database(
await llm_usage_recorder.record_usage_to_database(
model_info=model_info,
time_cost=time.time() - start_time,
model_usage=usage,
@@ -625,9 +625,9 @@ class LLMRequest:
logger.error(f"任务-'{task_name}' 模型-'{model_name}': 未知异常，错误信息-{str(e)}")
return -1, None # 不再重试请求该模型
@staticmethod
def _check_retry(
self,
remain_try: int,
remain_try: int,
retry_interval: int,
can_retry_msg: str,
cannot_retry_msg: str,
@@ -745,7 +745,8 @@ class LLMRequest:
)
return -1, None
def _build_tool_options(self, tools: Optional[List[Dict[str, Any]]]) -> Optional[List[ToolOption]]:
@staticmethod
def _build_tool_options(tools: Optional[List[Dict[str, Any]]]) -> Optional[List[ToolOption]]:
# sourcery skip: extract-method
"""构建工具选项列表"""
if not tools:
@@ -809,7 +810,8 @@ class LLMRequest:
return final_text
def _inject_random_noise(self, text: str, intensity: int) -> str:
@staticmethod
def _inject_random_noise(text: str, intensity: int) -> str:
"""在文本中注入随机乱码"""
import random
import string

78
src/main.py
View File

@@ -1,37 +1,35 @@
# 再用这个就写一行注释来混提交的我直接全部🌿飞😡
import asyncio
import time
import signal
import sys
import time
from maim_message import MessageServer
from src.common.remote import TelemetryHeartBeatTask
from src.manager.async_task_manager import async_task_manager
from src.chat.utils.statistic import OnlineTimeRecordTask, StatisticOutputTask
from src.common.remote import TelemetryHeartBeatTask
from src.chat.emoji_system.emoji_manager import get_emoji_manager
from src.chat.message_receive.chat_stream import get_chat_manager
from src.config.config import global_config
from src.chat.message_receive.bot import chat_bot
from src.common.logger import get_logger
from src.individuality.individuality import get_individuality, Individuality
from src.common.server import get_global_server, Server
from src.mood.mood_manager import mood_manager
from rich.traceback import install
from src.schedule.schedule_manager import schedule_manager
from src.schedule.monthly_plan_manager import monthly_plan_manager
from src.plugin_system.core.event_manager import event_manager
from src.plugin_system.base.component_types import EventType
# from src.api.main import start_api_server
# 导入新的插件管理器和热重载管理器
from src.plugin_system.core.plugin_manager import plugin_manager
from src.plugin_system.core.plugin_hot_reload import hot_reload_manager
from src.chat.emoji_system.emoji_manager import get_emoji_manager
from src.chat.memory_system.Hippocampus import hippocampus_manager
from src.chat.message_receive.bot import chat_bot
from src.chat.message_receive.chat_stream import get_chat_manager
from src.chat.utils.statistic import OnlineTimeRecordTask, StatisticOutputTask
from src.common.logger import get_logger
# 导入消息API和traceback模块
from src.common.message import get_global_api
from src.common.remote import TelemetryHeartBeatTask
from src.common.server import get_global_server, Server
from src.config.config import global_config
from src.individuality.individuality import get_individuality, Individuality
from src.manager.async_task_manager import async_task_manager
from src.mood.mood_manager import mood_manager
from src.plugin_system.base.component_types import EventType
from src.plugin_system.core.event_manager import event_manager
from src.plugin_system.core.plugin_hot_reload import hot_reload_manager
# 导入新的插件管理器和热重载管理器
from src.plugin_system.core.plugin_manager import plugin_manager
from src.schedule.monthly_plan_manager import monthly_plan_manager
from src.schedule.schedule_manager import schedule_manager
from src.chat.memory_system.Hippocampus import hippocampus_manager
# from src.api.main import start_api_server
if not global_config.memory.enable_memory:
import src.chat.memory_system.Hippocampus as hippocampus_module
@@ -40,7 +38,11 @@ if not global_config.memory.enable_memory:
def initialize(self):
pass
def get_hippocampus(self):
async def initialize_async(self):
pass
@staticmethod
def get_hippocampus():
return None
async def build_memory(self):
@@ -52,9 +54,9 @@ if not global_config.memory.enable_memory:
async def consolidate_memory(self):
pass
@staticmethod
async def get_memory_from_text(
self,
text: str,
text: str,
max_memory_num: int = 3,
max_memory_length: int = 2,
max_depth: int = 3,
@@ -62,20 +64,24 @@ if not global_config.memory.enable_memory:
) -> list:
return []
@staticmethod
async def get_memory_from_topic(
self, valid_keywords: list[str], max_memory_num: int = 3, max_memory_length: int = 2, max_depth: int = 3
valid_keywords: list[str], max_memory_num: int = 3, max_memory_length: int = 2, max_depth: int = 3
) -> list:
return []
@staticmethod
async def get_activate_from_text(
self, text: str, max_depth: int = 3, fast_retrieval: bool = False
text: str, max_depth: int = 3, fast_retrieval: bool = False
) -> tuple[float, list[str]]:
return 0.0, []
def get_memory_from_keyword(self, keyword: str, max_depth: int = 2) -> list:
@staticmethod
def get_memory_from_keyword(keyword: str, max_depth: int = 2) -> list:
return []
def get_all_node_names(self) -> list:
@staticmethod
def get_all_node_names() -> list:
return []
hippocampus_module.hippocampus_manager = MockHippocampusManager()
@@ -111,7 +117,8 @@ class MainSystem:
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
def _cleanup(self):
@staticmethod
def _cleanup():
"""清理资源"""
try:
# 停止消息重组器
@@ -248,7 +255,7 @@ MoFox_Bot(第三方修改版)
logger.info("聊天管理器初始化成功")
# 初始化记忆系统
self.hippocampus_manager.initialize()
await self.hippocampus_manager.initialize_async()
logger.info("记忆系统初始化成功")
# 初始化LPMM知识库
@@ -283,7 +290,7 @@ MoFox_Bot(第三方修改版)
if global_config.planning_system.monthly_plan_enable:
logger.info("正在初始化月度计划管理器...")
try:
await monthly_plan_manager.start_monthly_plan_generation()
await monthly_plan_manager.initialize()
logger.info("月度计划管理器初始化成功")
except Exception as e:
logger.error(f"月度计划管理器初始化失败: {e}")
@@ -291,8 +298,7 @@ MoFox_Bot(第三方修改版)
# 初始化日程管理器
if global_config.planning_system.schedule_enable:
logger.info("日程表功能已启用，正在初始化管理器...")
await schedule_manager.load_or_generate_today_schedule()
await schedule_manager.start_daily_schedule_generation()
await schedule_manager.initialize()
logger.info("日程表管理器初始化成功。")
try:

8
src/mais4u/mais4u_chat/body_emotion_action_manager.py
View File

@@ -118,14 +118,14 @@ class ChatAction:
self.regression_count = 0
message_time: float = message.message_info.time # type: ignore
message_list_before_now = get_raw_msg_by_timestamp_with_chat_inclusive(
message_list_before_now = await get_raw_msg_by_timestamp_with_chat_inclusive(
chat_id=self.chat_id,
timestamp_start=self.last_change_time,
timestamp_end=message_time,
limit=15,
limit_mode="last",
)
chat_talking_prompt = build_readable_messages(
chat_talking_prompt = await build_readable_messages(
message_list_before_now,
replace_bot_name=True,
merge_messages=False,
@@ -182,14 +182,14 @@ class ChatAction:
async def regress_action(self):
message_time = time.time()
message_list_before_now = get_raw_msg_by_timestamp_with_chat_inclusive(
message_list_before_now = await get_raw_msg_by_timestamp_with_chat_inclusive(
chat_id=self.chat_id,
timestamp_start=self.last_change_time,
timestamp_end=message_time,
limit=10,
limit_mode="last",
)
chat_talking_prompt = build_readable_messages(
chat_talking_prompt = await build_readable_messages(
message_list_before_now,
replace_bot_name=True,
merge_messages=False,

21
src/mais4u/mais4u_chat/s4u_chat.py
View File

@@ -58,7 +58,8 @@ class MessageSenderContainer:
"""恢复发送。"""
self._paused_event.set()
def _calculate_typing_delay(self, text: str) -> float:
@staticmethod
def _calculate_typing_delay(text: str) -> float:
"""根据文本长度计算模拟打字延迟。"""
chars_per_second = s4u_config.chars_per_second
min_delay = s4u_config.min_typing_delay
@@ -150,6 +151,10 @@ class MessageSenderContainer:
if self._task:
await self._task
@property
def task(self):
return self._task
class S4UChatManager:
def __init__(self):
@@ -177,6 +182,7 @@ class S4UChat:
def __init__(self, chat_stream: ChatStream):
"""初始化 S4UChat 实例。"""
self.last_msg_id = self.msg_id
self.chat_stream = chat_stream
self.stream_id = chat_stream.stream_id
self.stream_name = get_chat_manager().get_stream_name(self.stream_id) or self.stream_id
@@ -206,7 +212,8 @@ class S4UChat:
logger.info(f"[{self.stream_name}] S4UChat with two-queue system initialized.")
def _get_priority_info(self, message: MessageRecv) -> dict:
@staticmethod
def _get_priority_info(message: MessageRecv) -> dict:
"""安全地从消息中提取和解析 priority_info"""
priority_info_raw = message.priority_info
priority_info = {}
@@ -219,7 +226,8 @@ class S4UChat:
priority_info = priority_info_raw
return priority_info
def _is_vip(self, priority_info: dict) -> bool:
@staticmethod
def _is_vip(priority_info: dict) -> bool:
"""检查消息是否来自VIP用户。"""
return priority_info.get("message_type") == "vip"
@@ -468,7 +476,6 @@ class S4UChat:
await asyncio.sleep(1)
def get_processing_message_id(self):
self.last_msg_id = self.msg_id
self.msg_id = f"{time.time()}_{random.randint(1000, 9999)}"
async def _generate_and_send(self, message: MessageRecv):
@@ -565,7 +572,7 @@ class S4UChat:
# 确保发送器被妥善关闭（即使已关闭，再次调用也是安全的）
sender_container.resume()
if not sender_container._task.done():
if not sender_container.task.done():
await sender_container.close()
await sender_container.join()
logger.info(f"[{self.stream_name}] _generate_and_send 任务结束，资源已清理。")
@@ -586,3 +593,7 @@ class S4UChat:
await self._processing_task
except asyncio.CancelledError:
logger.info(f"处理任务已成功取消: {self.stream_name}")
@property
def new_message_event(self):
return self._new_message_event

11
src/mais4u/mais4u_chat/s4u_mood_manager.py
View File

@@ -124,7 +124,8 @@ class ChatMood:
# 发送初始情绪状态到ws端
asyncio.create_task(self.send_emotion_update(self.mood_values))
def _parse_numerical_mood(self, response: str) -> dict[str, int] | None:
@staticmethod
def _parse_numerical_mood(response: str) -> dict[str, int] | None:
try:
# The LLM might output markdown with json inside
if "```json" in response:
@@ -159,14 +160,14 @@ class ChatMood:
self.regression_count = 0
message_time: float = message.message_info.time # type: ignore
message_list_before_now = get_raw_msg_by_timestamp_with_chat_inclusive(
message_list_before_now = await get_raw_msg_by_timestamp_with_chat_inclusive(
chat_id=self.chat_id,
timestamp_start=self.last_change_time,
timestamp_end=message_time,
limit=10,
limit_mode="last",
)
chat_talking_prompt = build_readable_messages(
chat_talking_prompt = await build_readable_messages(
message_list_before_now,
replace_bot_name=True,
merge_messages=False,
@@ -238,14 +239,14 @@ class ChatMood:
async def regress_mood(self):
message_time = time.time()
message_list_before_now = get_raw_msg_by_timestamp_with_chat_inclusive(
message_list_before_now = await get_raw_msg_by_timestamp_with_chat_inclusive(
chat_id=self.chat_id,
timestamp_start=self.last_change_time,
timestamp_end=message_time,
limit=5,
limit_mode="last",
)
chat_talking_prompt = build_readable_messages(
chat_talking_prompt = await build_readable_messages(
message_list_before_now,
replace_bot_name=True,
merge_messages=False,

17
src/mais4u/mais4u_chat/s4u_msg_processor.py
View File

@@ -161,7 +161,8 @@ class S4UMessageProcessor:
else:
logger.info(f"[S4U]{userinfo.user_nickname}:{message.processed_plain_text}")
async def handle_internal_message(self, message: MessageRecvS4U):
@staticmethod
async def handle_internal_message(message: MessageRecvS4U):
if message.is_internal:
group_info = GroupInfo(platform="amaidesu_default", group_id=660154, group_name="内心")
@@ -173,7 +174,7 @@ class S4UMessageProcessor:
message.message_info.platform = s4u_chat.chat_stream.platform
s4u_chat.internal_message.append(message)
s4u_chat._new_message_event.set()
s4u_chat.new_message_event.set()
logger.info(
f"[{s4u_chat.stream_name}] 添加内部消息-------------------------------------------------------: {message.processed_plain_text}"
@@ -182,20 +183,23 @@ class S4UMessageProcessor:
return True
return False
async def handle_screen_message(self, message: MessageRecvS4U):
@staticmethod
async def handle_screen_message(message: MessageRecvS4U):
if message.is_screen:
screen_manager.set_screen(message.screen_info)
return True
return False
async def hadle_if_voice_done(self, message: MessageRecvS4U):
@staticmethod
async def hadle_if_voice_done(message: MessageRecvS4U):
if message.voice_done:
s4u_chat = get_s4u_chat_manager().get_or_create_chat(message.chat_stream)
s4u_chat.voice_done = message.voice_done
return True
return False
async def check_if_fake_gift(self, message: MessageRecvS4U) -> bool:
@staticmethod
async def check_if_fake_gift(message: MessageRecvS4U) -> bool:
"""检查消息是否为假礼物"""
if message.is_gift:
return False
@@ -227,7 +231,8 @@ class S4UMessageProcessor:
return True # 非礼物消息，继续正常处理
async def _handle_context_web_update(self, chat_id: str, message: MessageRecv):
@staticmethod
async def _handle_context_web_update(chat_id: str, message: MessageRecv):
"""处理上下文网页更新的独立task
Args:

24
src/mais4u/mais4u_chat/s4u_prompt.py
View File

@@ -98,7 +98,8 @@ class PromptBuilder:
self.prompt_built = ""
self.activate_messages = ""
async def build_expression_habits(self, chat_stream: ChatStream, chat_history, target):
@staticmethod
async def build_expression_habits(chat_stream: ChatStream, chat_history, target):
style_habits = []
grammar_habits = []
@@ -133,7 +134,8 @@ class PromptBuilder:
return expression_habits_block
async def build_relation_info(self, chat_stream) -> str:
@staticmethod
async def build_relation_info(chat_stream) -> str:
is_group_chat = bool(chat_stream.group_info)
who_chat_in_group = []
if is_group_chat:
@@ -167,7 +169,8 @@ class PromptBuilder:
)
return relation_prompt
async def build_memory_block(self, text: str) -> str:
@staticmethod
async def build_memory_block(text: str) -> str:
related_memory = await hippocampus_manager.get_memory_from_text(
text=text, max_memory_num=2, max_memory_length=2, max_depth=3, fast_retrieval=False
)
@@ -179,7 +182,8 @@ class PromptBuilder:
return await global_prompt_manager.format_prompt("memory_prompt", memory_info=related_memory_info)
return ""
def build_chat_history_prompts(self, chat_stream: ChatStream, message: MessageRecvS4U):
@staticmethod
async def build_chat_history_prompts(chat_stream: ChatStream, message: MessageRecvS4U):
message_list_before_now = get_raw_msg_before_timestamp_with_chat(
chat_id=chat_stream.stream_id,
timestamp=time.time(),
@@ -213,7 +217,7 @@ class PromptBuilder:
background_dialogue_prompt = ""
if background_dialogue_list:
context_msgs = background_dialogue_list[-s4u_config.max_context_message_length :]
background_dialogue_prompt_str = build_readable_messages(
background_dialogue_prompt_str = await build_readable_messages(
context_msgs,
timestamp_mode="normal_no_YMD",
show_pic=False,
@@ -262,7 +266,7 @@ class PromptBuilder:
timestamp=time.time(),
limit=20,
)
all_dialogue_prompt_str = build_readable_messages(
all_dialogue_prompt_str = await build_readable_messages(
all_dialogue_prompt,
timestamp_mode="normal_no_YMD",
show_pic=False,
@@ -270,7 +274,8 @@ class PromptBuilder:
return core_msg_str, background_dialogue_prompt, all_dialogue_prompt_str
def build_gift_info(self, message: MessageRecvS4U):
@staticmethod
def build_gift_info(message: MessageRecvS4U):
if message.is_gift:
return f"这是一条礼物信息，{message.gift_name} x{message.gift_count}，请注意这位用户"
else:
@@ -279,7 +284,8 @@ class PromptBuilder:
return ""
def build_sc_info(self, message: MessageRecvS4U):
@staticmethod
def build_sc_info(message: MessageRecvS4U):
super_chat_manager = get_super_chat_manager()
return super_chat_manager.build_superchat_summary_string(message.chat_stream.stream_id)
@@ -310,7 +316,7 @@ class PromptBuilder:
self.build_expression_habits(chat_stream, message_txt, sender_name),
)
core_dialogue_prompt, background_dialogue_prompt, all_dialogue_prompt = self.build_chat_history_prompts(
core_dialogue_prompt, background_dialogue_prompt, all_dialogue_prompt = await self.build_chat_history_prompts(
chat_stream, message
)

3
src/mais4u/mais4u_chat/s4u_stream_generator.py
View File

@@ -49,7 +49,8 @@ class S4UStreamGenerator:
self.chat_stream = None
async def build_last_internal_message(self, message: MessageRecvS4U, previous_reply_context: str = ""):
@staticmethod
async def build_last_internal_message(message: MessageRecvS4U, previous_reply_context: str = ""):
# person_id = PersonInfoManager.get_person_id(
# message.chat_stream.user_info.platform, message.chat_stream.user_info.user_id
# )

3
src/mais4u/mais4u_chat/super_chat_manager.py
View File

@@ -105,7 +105,8 @@ class SuperChatManager:
logger.error(f"清理过期SuperChat时出错: {e}", exc_info=True)
await asyncio.sleep(60) # 出错时等待更长时间
def _calculate_expire_time(self, price: float) -> float:
@staticmethod
def _calculate_expire_time(price: float) -> float:
"""根据SuperChat金额计算过期时间"""
current_time = time.time()

2
src/mais4u/s4u_config.py
View File

@@ -78,7 +78,7 @@ class S4UConfigBase:
except Exception as e:
raise RuntimeError(f"Failed to convert field '{field_name}' to target type: {e}") from e
return cls(**init_args)
return cls()
@classmethod
def _convert_field(cls, value: Any, field_type: Type[Any]) -> Any:

4
src/manager/async_task_manager.py
View File

@@ -1,4 +1,4 @@
from abc import abstractmethod
from abc import abstractmethod, ABCMeta
import asyncio
from asyncio import Task, Event, Lock
@@ -9,7 +9,7 @@ from src.common.logger import get_logger
logger = get_logger("async_task_manager")
class AsyncTask:
class AsyncTask(metaclass=ABCMeta):
"""异步任务基类"""
def __init__(self, task_name: str | None = None, wait_before_start: int = 0, run_interval: int = 0):

8
src/mood/mood_manager.py
View File

@@ -98,14 +98,14 @@ class ChatMood:
)
message_time: float = message.message_info.time # type: ignore
message_list_before_now = get_raw_msg_by_timestamp_with_chat_inclusive(
message_list_before_now = await get_raw_msg_by_timestamp_with_chat_inclusive(
chat_id=self.chat_id,
timestamp_start=self.last_change_time,
timestamp_end=message_time,
limit=int(global_config.chat.max_context_size / 3),
limit_mode="last",
)
chat_talking_prompt = build_readable_messages(
chat_talking_prompt = await build_readable_messages(
message_list_before_now,
replace_bot_name=True,
merge_messages=False,
@@ -147,14 +147,14 @@ class ChatMood:
async def regress_mood(self):
message_time = time.time()
message_list_before_now = get_raw_msg_by_timestamp_with_chat_inclusive(
message_list_before_now = await get_raw_msg_by_timestamp_with_chat_inclusive(
chat_id=self.chat_id,
timestamp_start=self.last_change_time,
timestamp_end=message_time,
limit=15,
limit_mode="last",
)
chat_talking_prompt = build_readable_messages(
chat_talking_prompt = await build_readable_messages(
message_list_before_now,
replace_bot_name=True,
merge_messages=False,

305
src/person_info/person_info.py
View File

@@ -1,18 +1,18 @@
import copy
import hashlib
import datetime
import asyncio
import orjson
import hashlib
import time
from json_repair import repair_json
from typing import Any, Callable, Dict, Union, Optional
import orjson
from json_repair import repair_json
from sqlalchemy import select
from src.common.logger import get_logger
from src.common.database.sqlalchemy_models import PersonInfo
from src.common.database.sqlalchemy_database_api import get_db_session
from src.llm_models.utils_model import LLMRequest
from src.common.database.sqlalchemy_models import PersonInfo
from src.common.logger import get_logger
from src.config.config import global_config, model_config
from src.llm_models.utils_model import LLMRequest
"""
PersonInfoManager 类方法功能摘要：
@@ -73,14 +73,15 @@ class PersonInfoManager:
# # 初始化时读取所有person_name
try:
pass
# 在这里获取会话
with get_db_session() as session:
for record in session.execute(
select(PersonInfo.person_id, PersonInfo.person_name).where(PersonInfo.person_name.is_not(None))
).fetchall():
if record.person_name:
self.person_name_list[record.person_id] = record.person_name
logger.debug(f"已加载 {len(self.person_name_list)} 个用户名称 (SQLAlchemy)")
# with get_db_session() as session:
# for record in session.execute(
# select(PersonInfo.person_id, PersonInfo.person_name).where(PersonInfo.person_name.is_not(None))
# ).fetchall():
# if record.person_name:
# self.person_name_list[record.person_id] = record.person_name
# logger.debug(f"已加载 {len(self.person_name_list)} 个用户名称 (SQLAlchemy)")
except Exception as e:
logger.error(f"从 SQLAlchemy 加载 person_name_list 失败: {e}")
@@ -102,23 +103,26 @@ class PersonInfoManager:
"""判断是否认识某人"""
person_id = self.get_person_id(platform, user_id)
def _db_check_known_sync(p_id: str):
async def _db_check_known_async(p_id: str):
# 在需要时获取会话
with get_db_session() as session:
return session.execute(select(PersonInfo).where(PersonInfo.person_id == p_id)).scalar() is not None
async with get_db_session() as session:
return (
await session.execute(select(PersonInfo).where(PersonInfo.person_id == p_id))
).scalar() is not None
try:
return await asyncio.to_thread(_db_check_known_sync, person_id)
return await _db_check_known_async(person_id)
except Exception as e:
logger.error(f"检查用户 {person_id} 是否已知时出错 (SQLAlchemy): {e}")
return False
def get_person_id_by_person_name(self, person_name: str) -> str:
@staticmethod
async def get_person_id_by_person_name(person_name: str) -> str:
"""根据用户名获取用户ID"""
try:
# 在需要时获取会话
with get_db_session() as session:
record = session.execute(select(PersonInfo).where(PersonInfo.person_name == person_name)).scalar()
async with get_db_session() as session:
record = (await session.execute(select(PersonInfo).where(PersonInfo.person_name == person_name))).scalar()
return record.person_id if record else ""
except Exception as e:
logger.error(f"根据用户名 {person_name} 获取用户ID时出错 (SQLAlchemy): {e}")
@@ -172,21 +176,21 @@ class PersonInfoManager:
final_data[key] = orjson.dumps([]).decode("utf-8")
# If it's already a string, assume it's valid JSON or a non-JSON string field
def _db_create_sync(p_data: dict):
with get_db_session() as session:
async def _db_create_async(p_data: dict):
async with get_db_session() as session:
try:
new_person = PersonInfo(**p_data)
session.add(new_person)
session.commit()
await session.commit()
return True
except Exception as e:
logger.error(f"创建 PersonInfo 记录 {p_data.get('person_id')} 失败 (SQLAlchemy): {e}")
return False
await asyncio.to_thread(_db_create_sync, final_data)
await _db_create_async(final_data)
async def _safe_create_person_info(self, person_id: str, data: Optional[dict] = None):
@staticmethod
async def _safe_create_person_info(person_id: str, data: Optional[dict] = None):
"""安全地创建用户信息，处理竞态条件"""
if not person_id:
logger.debug("创建失败，person_id不存在")
@@ -229,11 +233,11 @@ class PersonInfoManager:
elif final_data[key] is None: # Default for lists is [], store as "[]"
final_data[key] = orjson.dumps([]).decode("utf-8")
def _db_safe_create_sync(p_data: dict):
with get_db_session() as session:
async def _db_safe_create_async(p_data: dict):
async with get_db_session() as session:
try:
existing = session.execute(
select(PersonInfo).where(PersonInfo.person_id == p_data["person_id"])
existing = (
await session.execute(select(PersonInfo).where(PersonInfo.person_id == p_data["person_id"]))
).scalar()
if existing:
logger.debug(f"用户 {p_data['person_id']} 已存在，跳过创建")
@@ -242,18 +246,17 @@ class PersonInfoManager:
# 尝试创建
new_person = PersonInfo(**p_data)
session.add(new_person)
session.commit()
await session.commit()
return True
except Exception as e:
if "UNIQUE constraint failed" in str(e):
logger.debug(f"检测到并发创建用户 {p_data.get('person_id')}，跳过错误")
return True # 其他协程已创建，视为成功
return True
else:
logger.error(f"创建 PersonInfo 记录 {p_data.get('person_id')} 失败 (SQLAlchemy): {e}")
return False
await asyncio.to_thread(_db_safe_create_sync, final_data)
await _db_safe_create_async(final_data)
async def update_one_field(self, person_id: str, field_name: str, value, data: Optional[Dict] = None):
"""更新某一个字段，会补全"""
@@ -270,37 +273,33 @@ class PersonInfoManager:
elif value is None: # Store None as "[]" for JSON list fields
processed_value = orjson.dumps([]).decode("utf-8")
def _db_update_sync(p_id: str, f_name: str, val_to_set):
async def _db_update_async(p_id: str, f_name: str, val_to_set):
start_time = time.time()
with get_db_session() as session:
async with get_db_session() as session:
try:
record = session.execute(select(PersonInfo).where(PersonInfo.person_id == p_id)).scalar()
record = (await session.execute(select(PersonInfo).where(PersonInfo.person_id == p_id))).scalar()
query_time = time.time()
if record:
setattr(record, f_name, val_to_set)
save_time = time.time()
total_time = save_time - start_time
if total_time > 0.5: # 如果超过500ms就记录日志
if total_time > 0.5:
logger.warning(
f"数据库更新操作耗时 {total_time:.3f}秒 (查询: {query_time - start_time:.3f}s, 保存: {save_time - query_time:.3f}s) person_id={p_id}, field={f_name}"
)
session.commit()
return True, False # Found and updated, no creation needed
await session.commit()
return True, False
else:
total_time = time.time() - start_time
if total_time > 0.5:
logger.warning(f"数据库查询操作耗时 {total_time:.3f}秒 person_id={p_id}, field={f_name}")
return False, True # Not found, needs creation
return False, True
except Exception as e:
total_time = time.time() - start_time
logger.error(f"数据库操作异常，耗时 {total_time:.3f}秒: {e}")
raise
found, needs_creation = await asyncio.to_thread(_db_update_sync, person_id, field_name, processed_value)
found, needs_creation = await _db_update_async(person_id, field_name, processed_value)
if needs_creation:
logger.info(f"{person_id} 不存在，将新建。")
@@ -338,13 +337,13 @@ class PersonInfoManager:
logger.debug(f"检查字段'{field_name}'失败，未在 PersonInfo SQLAlchemy 模型中定义。")
return False
def _db_has_field_sync(p_id: str, f_name: str):
with get_db_session() as session:
record = session.execute(select(PersonInfo).where(PersonInfo.person_id == p_id)).scalar()
async def _db_has_field_async(p_id: str, f_name: str):
async with get_db_session() as session:
record = (await session.execute(select(PersonInfo).where(PersonInfo.person_id == p_id))).scalar()
return bool(record)
try:
return await asyncio.to_thread(_db_has_field_sync, person_id, field_name)
return await _db_has_field_async(person_id, field_name)
except Exception as e:
logger.error(f"检查字段 {field_name} for {person_id} 时出错 (SQLAlchemy): {e}")
return False
@@ -449,14 +448,14 @@ class PersonInfoManager:
logger.info(f"尝试给用户{user_nickname} {person_id} 取名，但是 {generated_nickname} 已存在，重试中...")
else:
def _db_check_name_exists_sync(name_to_check):
with get_db_session() as session:
async def _db_check_name_exists_async(name_to_check):
async with get_db_session() as session:
return (
session.execute(select(PersonInfo).where(PersonInfo.person_name == name_to_check)).scalar()
(await session.execute(select(PersonInfo).where(PersonInfo.person_name == name_to_check))).scalar()
is not None
)
if await asyncio.to_thread(_db_check_name_exists_sync, generated_nickname):
if await _db_check_name_exists_async(generated_nickname):
is_duplicate = True
current_name_set.add(generated_nickname)
@@ -492,91 +491,65 @@ class PersonInfoManager:
logger.debug("删除失败：person_id 不能为空")
return
def _db_delete_sync(p_id: str):
async def _db_delete_async(p_id: str):
try:
with get_db_session() as session:
record = session.execute(select(PersonInfo).where(PersonInfo.person_id == p_id)).scalar()
async with get_db_session() as session:
record = (await session.execute(select(PersonInfo).where(PersonInfo.person_id == p_id))).scalar()
if record:
session.delete(record)
session.commit()
return 1
await session.delete(record)
await session.commit()
return 1
return 0
except Exception as e:
logger.error(f"删除 PersonInfo {p_id} 失败 (SQLAlchemy): {e}")
return 0
deleted_count = await asyncio.to_thread(_db_delete_sync, person_id)
deleted_count = await _db_delete_async(person_id)
if deleted_count > 0:
logger.debug(f"删除成功：person_id={person_id} (Peewee)")
logger.debug(f"删除成功：person_id={person_id}")
else:
logger.debug(f"删除失败：未找到 person_id={person_id} 或删除未影响行 (Peewee)")
logger.debug(f"删除失败：未找到 person_id={person_id} 或删除未影响行")
@staticmethod
async def get_value(person_id: str, field_name: str):
"""获取指定用户指定字段的值"""
default_value_for_field = person_info_default.get(field_name)
if field_name in JSON_SERIALIZED_FIELDS and default_value_for_field is None:
default_value_for_field = [] # Ensure JSON fields default to [] if not in DB
def get_value(person_id: str, field_name: str) -> Any:
"""获取单个字段值（同步版本）"""
if not person_id:
logger.debug("get_value获取失败：person_id不能为空")
return None
def _db_get_value_sync(p_id: str, f_name: str):
with get_db_session() as session:
record = session.execute(select(PersonInfo).where(PersonInfo.person_id == p_id)).scalar()
if record:
val = getattr(record, f_name, None)
if f_name in JSON_SERIALIZED_FIELDS:
if isinstance(val, str):
try:
return orjson.loads(val)
except orjson.JSONDecodeError:
logger.warning(f"字段 {f_name} for {p_id} 包含无效JSON: {val}. 返回默认值.")
return [] # Default for JSON fields on error
elif val is None: # Field exists in DB but is None
return [] # Default for JSON fields
# If val is already a list/dict (e.g. if somehow set without serialization)
return val # Should ideally not happen if update_one_field is always used
return val
return None # Record not found
import asyncio
async def _get_record_sync():
async with get_db_session() as session:
return (await session.execute(select(PersonInfo).where(PersonInfo.person_id == person_id))).scalar()
try:
value_from_db = await asyncio.to_thread(_db_get_value_sync, person_id, field_name)
if value_from_db is not None:
return value_from_db
record = asyncio.run(_get_record_sync())
except RuntimeError:
# 如果当前线程已经有事件循环在运行，则使用现有的循环
loop = asyncio.get_running_loop()
record = loop.run_until_complete(_get_record_sync())
model_fields = [column.name for column in PersonInfo.__table__.columns]
if field_name not in model_fields:
if field_name in person_info_default:
return default_value_for_field
logger.warning(f"字段 {field_name} 在 person_info_default 中未定义，且在数据库中未找到。")
return None # Ultimate fallback
except Exception as e:
logger.error(f"获取字段 {field_name} for {person_id} 时出错 (Peewee): {e}")
# Fallback to default in case of any error during DB access
return default_value_for_field if field_name in person_info_default else None
logger.debug(f"字段'{field_name}'不在SQLAlchemy模型中，使用默认配置值。")
return copy.deepcopy(person_info_default[field_name])
else:
logger.debug(f"get_value查询失败：字段'{field_name}'未在SQLAlchemy模型和默认配置中定义。")
return None
@staticmethod
def get_value_sync(person_id: str, field_name: str):
"""同步获取指定用户指定字段的值"""
default_value_for_field = person_info_default.get(field_name)
with get_db_session() as session:
if field_name in JSON_SERIALIZED_FIELDS and default_value_for_field is None:
default_value_for_field = []
if record := session.execute(select(PersonInfo).where(PersonInfo.person_id == person_id)).scalar():
val = getattr(record, field_name, None)
if field_name in JSON_SERIALIZED_FIELDS:
if isinstance(val, str):
try:
return orjson.loads(val)
except orjson.JSONDecodeError:
logger.warning(f"字段 {field_name} for {person_id} 包含无效JSON: {val}. 返回默认值.")
return []
elif val is None:
return []
return val
return val
if field_name in person_info_default:
return default_value_for_field
logger.warning(f"字段 {field_name} 在 person_info_default 中未定义，且在数据库中未找到。")
return None
if record:
value = getattr(record, field_name)
if value is not None:
return value
else:
return copy.deepcopy(person_info_default.get(field_name))
else:
return copy.deepcopy(person_info_default.get(field_name))
@staticmethod
async def get_values(person_id: str, field_names: list) -> dict:
@@ -587,11 +560,11 @@ class PersonInfoManager:
result = {}
def _db_get_record_sync(p_id: str):
with get_db_session() as session:
return session.execute(select(PersonInfo).where(PersonInfo.person_id == p_id)).scalar()
async def _db_get_record_async(p_id: str):
async with get_db_session() as session:
return (await session.execute(select(PersonInfo).where(PersonInfo.person_id == p_id))).scalar()
record = await asyncio.to_thread(_db_get_record_sync, person_id)
record = await _db_get_record_async(person_id)
# 获取 SQLAlchemy 模型的所有字段名
model_fields = [column.name for column in PersonInfo.__table__.columns]
@@ -616,7 +589,6 @@ class PersonInfoManager:
result[field_name] = copy.deepcopy(person_info_default.get(field_name))
return result
@staticmethod
async def get_specific_value_list(
field_name: str,
@@ -628,14 +600,15 @@ class PersonInfoManager:
# 获取 SQLAlchemy 模型的所有字段名
model_fields = [column.name for column in PersonInfo.__table__.columns]
if field_name not in model_fields:
logger.error(f"字段检查失败：'{field_name}'未在 PersonInfo SQLAlchemy 模 modelo中定义")
logger.error(f"字段检查失败：'{field_name}'未在 PersonInfo SQLAlchemy 模型中定义")
return {}
def _db_get_specific_sync(f_name: str):
async def _db_get_specific_async(f_name: str):
found_results = {}
try:
with get_db_session() as session:
for record in session.execute(select(PersonInfo.person_id, getattr(PersonInfo, f_name))).fetchall():
async with get_db_session() as session:
result = await session.execute(select(PersonInfo.person_id, getattr(PersonInfo, f_name)))
for record in result.fetchall():
value = getattr(record, f_name)
if way(value):
found_results[record.person_id] = value
@@ -646,9 +619,9 @@ class PersonInfoManager:
return found_results
try:
return await asyncio.to_thread(_db_get_specific_sync, field_name)
return await _db_get_specific_async(field_name)
except Exception as e:
logger.error(f"执行 get_specific_value_list 线程时出错: {str(e)}", exc_info=True)
logger.error(f"执行 get_specific_value_list 时出错: {str(e)}", exc_info=True)
return {}
async def get_or_create_person(
@@ -661,40 +634,38 @@ class PersonInfoManager:
"""
person_id = self.get_person_id(platform, user_id)
def _db_get_or_create_sync(p_id: str, init_data: dict):
async def _db_get_or_create_async(p_id: str, init_data: dict):
"""原子性的获取或创建操作"""
with get_db_session() as session:
async with get_db_session() as session:
# 首先尝试获取现有记录
record = session.execute(select(PersonInfo).where(PersonInfo.person_id == p_id)).scalar()
record = (await session.execute(select(PersonInfo).where(PersonInfo.person_id == p_id))).scalar()
if record:
return record, False # 记录存在，未创建
# 记录不存在，尝试创建
try:
new_person = PersonInfo(**init_data)
session.add(new_person)
session.commit()
return session.execute(
select(PersonInfo).where(PersonInfo.person_id == p_id)
).scalar(), True # 创建成功
except Exception as e:
# 如果创建失败（可能是因为竞态条件），再次尝试获取
if "UNIQUE constraint failed" in str(e):
logger.debug(f"检测到并发创建用户 {p_id}，获取现有记录")
record = session.execute(select(PersonInfo).where(PersonInfo.person_id == p_id)).scalar()
if record:
return record, False # 其他协程已创建，返回现有记录
# 如果仍然失败，重新抛出异常
raise e
# 记录不存在，尝试创建
try:
new_person = PersonInfo(**init_data)
session.add(new_person)
await session.commit()
await session.refresh(new_person)
return new_person, True # 创建成功
except Exception as e:
# 如果创建失败（可能是因为竞态条件），再次尝试获取
if "UNIQUE constraint failed" in str(e):
logger.debug(f"检测到并发创建用户 {p_id}，获取现有记录")
record = (await session.execute(select(PersonInfo).where(PersonInfo.person_id == p_id))).scalar()
if record:
return record, False # 其他协程已创建，返回现有记录
# 如果仍然失败，重新抛出异常
raise e
unique_nickname = await self._generate_unique_person_name(nickname)
initial_data = {
"person_id": person_id,
"platform": platform,
"user_id": str(user_id),
"nickname": nickname,
"person_name": unique_nickname, # 使用群昵称作为person_name
"person_name": unique_nickname,
"name_reason": "从群昵称获取",
"know_times": 0,
"know_since": int(datetime.datetime.now().timestamp()),
@@ -704,7 +675,6 @@ class PersonInfoManager:
"forgotten_points": [],
}
# 序列化JSON字段
for key in JSON_SERIALIZED_FIELDS:
if key in initial_data:
if isinstance(initial_data[key], (list, dict)):
@@ -712,15 +682,14 @@ class PersonInfoManager:
elif initial_data[key] is None:
initial_data[key] = orjson.dumps([]).decode("utf-8")
# 获取 SQLAlchemy 模odel的所有字段名
model_fields = [column.name for column in PersonInfo.__table__.columns]
filtered_initial_data = {k: v for k, v in initial_data.items() if v is not None and k in model_fields}
record, was_created = await asyncio.to_thread(_db_get_or_create_sync, person_id, filtered_initial_data)
record, was_created = await _db_get_or_create_async(person_id, filtered_initial_data)
if was_created:
logger.info(f"用户 {platform}:{user_id} (person_id: {person_id}) 不存在，将创建新记录 (Peewee)。")
logger.info(f"已为 {person_id} 创建新记录，初始数据 (filtered for model): {filtered_initial_data}")
logger.info(f"用户 {platform}:{user_id} (person_id: {person_id}) 不存在，将创建新记录。")
logger.info(f"已为 {person_id} 创建新记录，初始数据: {filtered_initial_data}")
else:
logger.debug(f"用户 {platform}:{user_id} (person_id: {person_id}) 已存在，返回现有记录。")
@@ -740,11 +709,13 @@ class PersonInfoManager:
if not found_person_id:
def _db_find_by_name_sync(p_name_to_find: str):
with get_db_session() as session:
return session.execute(select(PersonInfo).where(PersonInfo.person_name == p_name_to_find)).scalar()
async def _db_find_by_name_async(p_name_to_find: str):
async with get_db_session() as session:
return (
await session.execute(select(PersonInfo).where(PersonInfo.person_name == p_name_to_find))
).scalar()
record = await asyncio.to_thread(_db_find_by_name_sync, person_name)
record = await _db_find_by_name_async(person_name)
if record:
found_person_id = record.person_id
if (

69
src/person_info/relationship_builder.py
View File

@@ -113,7 +113,7 @@ class RelationshipBuilder:
# 负责跟踪用户消息活动、管理消息段、清理过期数据
# ================================
def _update_message_segments(self, person_id: str, message_time: float):
async def _update_message_segments(self, person_id: str, message_time: float):
"""更新用户的消息段
Args:
@@ -126,11 +126,8 @@ class RelationshipBuilder:
segments = self.person_engaged_cache[person_id]
# 获取该消息前5条消息的时间作为潜在的开始时间
before_messages = get_raw_msg_before_timestamp_with_chat(self.chat_id, message_time, limit=5)
if before_messages:
potential_start_time = before_messages[0]["time"]
else:
potential_start_time = message_time
before_messages = await get_raw_msg_before_timestamp_with_chat(self.chat_id, message_time, limit=5)
potential_start_time = before_messages[0]["time"] if before_messages else message_time
# 如果没有现有消息段，创建新的
if not segments:
@@ -138,11 +135,10 @@ class RelationshipBuilder:
"start_time": potential_start_time,
"end_time": message_time,
"last_msg_time": message_time,
"message_count": self._count_messages_in_timerange(potential_start_time, message_time),
"message_count": await self._count_messages_in_timerange(potential_start_time, message_time),
}
segments.append(new_segment)
person_name = get_person_info_manager().get_value_sync(person_id, "person_name") or person_id
person_name = get_person_info_manager().get_value(person_id, "person_name") or person_id
logger.debug(
f"{self.log_prefix} 眼熟用户 {person_name} 在 {time.strftime('%H:%M:%S', time.localtime(potential_start_time))} - {time.strftime('%H:%M:%S', time.localtime(message_time))} 之间有 {new_segment['message_count']} 条消息"
)
@@ -153,57 +149,50 @@ class RelationshipBuilder:
last_segment = segments[-1]
# 计算从最后一条消息到当前消息之间的消息数量（不包含边界）
messages_between = self._count_messages_between(last_segment["last_msg_time"], message_time)
messages_between = await self._count_messages_between(last_segment["last_msg_time"], message_time)
if messages_between <= 10:
# 在10条消息内，延伸当前消息段
last_segment["end_time"] = message_time
last_segment["last_msg_time"] = message_time
# 重新计算整个消息段的消息数量
last_segment["message_count"] = self._count_messages_in_timerange(
last_segment["message_count"] = await self._count_messages_in_timerange(
last_segment["start_time"], last_segment["end_time"]
)
logger.debug(f"{self.log_prefix} 延伸用户 {person_id} 的消息段: {last_segment}")
else:
# 超过10条消息，结束当前消息段并创建新的
# 结束当前消息段：延伸到原消息段最后一条消息后5条消息的时间
current_time = time.time()
after_messages = get_raw_msg_by_timestamp_with_chat(
after_messages = await get_raw_msg_by_timestamp_with_chat(
self.chat_id, last_segment["last_msg_time"], current_time, limit=5, limit_mode="earliest"
)
if after_messages and len(after_messages) >= 5:
# 如果有足够的后续消息，使用第5条消息的时间作为结束时间
last_segment["end_time"] = after_messages[4]["time"]
# 重新计算当前消息段的消息数量
last_segment["message_count"] = self._count_messages_in_timerange(
last_segment["message_count"] = await self._count_messages_in_timerange(
last_segment["start_time"], last_segment["end_time"]
)
# 创建新的消息段
new_segment = {
"start_time": potential_start_time,
"end_time": message_time,
"last_msg_time": message_time,
"message_count": self._count_messages_in_timerange(potential_start_time, message_time),
"message_count": await self._count_messages_in_timerange(potential_start_time, message_time),
}
segments.append(new_segment)
person_info_manager = get_person_info_manager()
person_name = person_info_manager.get_value_sync(person_id, "person_name") or person_id
person_name = person_info_manager.get_value(person_id, "person_name") or person_id
logger.debug(
f"{self.log_prefix} 重新眼熟用户 {person_name} 创建新消息段（超过10条消息间隔）: {new_segment}"
)
self._save_cache()
def _count_messages_in_timerange(self, start_time: float, end_time: float) -> int:
async def _count_messages_in_timerange(self, start_time: float, end_time: float) -> int:
"""计算指定时间范围内的消息数量（包含边界）"""
messages = get_raw_msg_by_timestamp_with_chat_inclusive(self.chat_id, start_time, end_time)
messages = await get_raw_msg_by_timestamp_with_chat_inclusive(self.chat_id, start_time, end_time)
return len(messages)
def _count_messages_between(self, start_time: float, end_time: float) -> int:
async def _count_messages_between(self, start_time: float, end_time: float) -> int:
"""计算两个时间点之间的消息数量（不包含边界），用于间隔检查"""
return num_new_messages_since(self.chat_id, start_time, end_time)
return await num_new_messages_since(self.chat_id, start_time, end_time)
def _get_total_message_count(self, person_id: str) -> int:
"""获取用户所有消息段的总消息数量"""
@@ -314,18 +303,12 @@ class RelationshipBuilder:
if not self.person_engaged_cache:
return f"{self.log_prefix} 关系缓存为空"
status_lines = [f"{self.log_prefix} 关系缓存状态："]
status_lines.append(
f"最后处理消息时间：{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(self.last_processed_message_time)) if self.last_processed_message_time > 0 else '未设置'}"
)
status_lines.append(
f"最后清理时间：{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(self.last_cleanup_time)) if self.last_cleanup_time > 0 else '未执行'}"
)
status_lines.append(f"总用户数：{len(self.person_engaged_cache)}")
status_lines.append(
f"清理配置：{'启用' if SEGMENT_CLEANUP_CONFIG['enable_cleanup'] else '禁用'} (最大保存{SEGMENT_CLEANUP_CONFIG['max_segment_age_days']}天, 每用户最多{SEGMENT_CLEANUP_CONFIG['max_segments_per_user']}段)"
)
status_lines.append("")
status_lines = [f"{self.log_prefix} 关系缓存状态：",
f"最后处理消息时间：{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(self.last_processed_message_time)) if self.last_processed_message_time > 0 else '未设置'}",
f"最后清理时间：{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(self.last_cleanup_time)) if self.last_cleanup_time > 0 else '未执行'}",
f"总用户数：{len(self.person_engaged_cache)}",
f"清理配置：{'启用' if SEGMENT_CLEANUP_CONFIG['enable_cleanup'] else '禁用'} (最大保存{SEGMENT_CLEANUP_CONFIG['max_segment_age_days']}天, 每用户最多{SEGMENT_CLEANUP_CONFIG['max_segments_per_user']}段)",
""]
for person_id, segments in self.person_engaged_cache.items():
total_count = self._get_total_message_count(person_id)
@@ -356,7 +339,7 @@ class RelationshipBuilder:
self._cleanup_old_segments()
current_time = time.time()
if latest_messages := get_raw_msg_by_timestamp_with_chat(
if latest_messages := await get_raw_msg_by_timestamp_with_chat(
self.chat_id,
self.last_processed_message_time,
current_time,
@@ -375,7 +358,7 @@ class RelationshipBuilder:
and msg_time > self.last_processed_message_time
):
person_id = PersonInfoManager.get_person_id(platform, user_id)
self._update_message_segments(person_id, msg_time)
await self._update_message_segments(person_id, msg_time)
logger.debug(
f"{self.log_prefix} 更新用户 {person_id} 的消息段，消息时间：{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(msg_time))}"
)
@@ -385,8 +368,8 @@ class RelationshipBuilder:
users_to_build_relationship = []
for person_id, segments in self.person_engaged_cache.items():
total_message_count = self._get_total_message_count(person_id)
person_name = get_person_info_manager().get_value_sync(person_id, "person_name") or person_id
person_name = get_person_info_manager().get_value(person_id, "person_name") or person_id
if total_message_count >= max_build_threshold or (
total_message_count >= 5 and (immediate_build == person_id or immediate_build == "all")
):
@@ -445,7 +428,7 @@ class RelationshipBuilder:
start_date = time.strftime("%Y-%m-%d %H:%M", time.localtime(start_time))
# 获取该段的消息（包含边界）
segment_messages = get_raw_msg_by_timestamp_with_chat_inclusive(self.chat_id, start_time, end_time)
segment_messages = await get_raw_msg_by_timestamp_with_chat_inclusive(self.chat_id, start_time, end_time)
logger.debug(
f"消息段: {start_date} - {time.strftime('%Y-%m-%d %H:%M', time.localtime(end_time))}, 消息数: {len(segment_messages)}"
)

32
src/person_info/relationship_fetcher.py
View File

@@ -99,16 +99,22 @@ class RelationshipFetcher:
self._cleanup_expired_cache()
person_info_manager = get_person_info_manager()
person_name = await person_info_manager.get_value(person_id, "person_name")
short_impression = await person_info_manager.get_value(person_id, "short_impression")
nickname_str = await person_info_manager.get_value(person_id, "nickname")
platform = await person_info_manager.get_value(person_id, "platform")
person_info = await person_info_manager.get_values(
person_id, ["person_name", "short_impression", "nickname", "platform", "points"]
)
person_name = person_info.get("person_name")
short_impression = person_info.get("short_impression")
nickname_str = person_info.get("nickname")
platform = person_info.get("platform")
if person_name == nickname_str and not short_impression:
return ""
current_points = await person_info_manager.get_value(person_id, "points") or []
current_points = person_info.get("points")
if isinstance(current_points, str):
current_points = orjson.loads(current_points)
else:
current_points = current_points or []
# 按时间排序forgotten_points
current_points.sort(key=lambda x: x[2])
@@ -170,7 +176,8 @@ class RelationshipFetcher:
nickname_str = ",".join(global_config.bot.alias_names)
name_block = f"你的名字是{global_config.bot.nickname},你的昵称有{nickname_str}，有人也会用这些昵称称呼你。"
person_info_manager = get_person_info_manager()
person_name: str = await person_info_manager.get_value(person_id, "person_name") # type: ignore
person_info = await person_info_manager.get_values(person_id, ["person_name"])
person_name: str = person_info.get("person_name") # type: ignore
info_cache_block = self._build_info_cache_block()
@@ -252,7 +259,8 @@ class RelationshipFetcher:
person_info_manager = get_person_info_manager()
# 首先检查 info_list 缓存
info_list = await person_info_manager.get_value(person_id, "info_list") or []
person_info = await person_info_manager.get_values(person_id, ["info_list"])
info_list = person_info.get("info_list") or []
cached_info = None
# 查找对应的 info_type
@@ -279,8 +287,9 @@ class RelationshipFetcher:
# 如果缓存中没有，尝试从用户档案中提取
try:
person_impression = await person_info_manager.get_value(person_id, "impression")
points = await person_info_manager.get_value(person_id, "points")
person_info = await person_info_manager.get_values(person_id, ["impression", "points"])
person_impression = person_info.get("impression")
points = person_info.get("points")
# 构建印象信息块
if person_impression:
@@ -372,7 +381,8 @@ class RelationshipFetcher:
person_info_manager = get_person_info_manager()
# 获取现有的 info_list
info_list = await person_info_manager.get_value(person_id, "info_list") or []
person_info = await person_info_manager.get_values(person_id, ["info_list"])
info_list = person_info.get("info_list") or []
# 查找是否已存在相同 info_type 的记录
found_index = -1

11
src/person_info/relationship_manager.py
View File

@@ -118,7 +118,7 @@ class RelationshipManager:
name_mapping[replace_person_name] = f"用户{current_user}{user_count if user_count > 1 else ''}"
current_user = chr(ord(current_user) + 1)
readable_messages = build_readable_messages(
readable_messages = await build_readable_messages(
messages=user_messages, replace_bot_name=True, timestamp_mode="normal_no_YMD", truncate=True
)
@@ -492,7 +492,8 @@ class RelationshipManager:
return current_points
def calculate_time_weight(self, point_time: str, current_time: str) -> float:
@staticmethod
def calculate_time_weight(point_time: str, current_time: str) -> float:
"""计算基于时间的权重系数"""
try:
point_timestamp = datetime.strptime(point_time, "%Y-%m-%d %H:%M:%S")
@@ -516,7 +517,8 @@ class RelationshipManager:
logger.error(f"计算时间权重失败: {e}")
return 0.5 # 发生错误时返回中等权重
def tfidf_similarity(self, s1, s2):
@staticmethod
def tfidf_similarity(s1, s2):
"""
使用 TF-IDF 和余弦相似度计算两个句子的相似性。
"""
@@ -551,7 +553,8 @@ class RelationshipManager:
# 返回 s1 和 s2 的相似度
return similarity_matrix[0, 1]
def sequence_similarity(self, s1, s2):
@staticmethod
def sequence_similarity(s1, s2):
"""
使用 SequenceMatcher 计算两个句子的相似性。
"""

2
src/plugin_system/apis/__init__.py
View File

@@ -19,6 +19,7 @@ from src.plugin_system.apis import (
send_api,
tool_api,
permission_api,
schedule_api
)
from src.plugin_system.apis.chat_api import ChatManager as context_api
from .logging_api import get_logger
@@ -42,4 +43,5 @@ __all__ = [
"tool_api",
"permission_api",
"context_api",
"schedule_api",
]

12
src/plugin_system/apis/cross_context_api.py
View File

@@ -53,14 +53,14 @@ async def build_cross_context_normal(chat_stream: ChatStream, other_chat_infos:
continue
try:
messages = get_raw_msg_before_timestamp_with_chat(
messages = await get_raw_msg_before_timestamp_with_chat(
chat_id=stream_id,
timestamp=time.time(),
limit=5, # 可配置
)
if messages:
chat_name = get_chat_manager().get_stream_name(stream_id) or chat_raw_id
formatted_messages, _ = build_readable_messages_with_id(messages, timestamp_mode="relative")
formatted_messages, _ = await build_readable_messages_with_id(messages, timestamp_mode="relative")
cross_context_messages.append(f'[以下是来自"{chat_name}"的近期消息]\n{formatted_messages}')
except Exception as e:
logger.error(f"获取聊天 {chat_raw_id} 的消息失败: {e}")
@@ -92,7 +92,7 @@ async def build_cross_context_s4u(
continue
try:
messages = get_raw_msg_before_timestamp_with_chat(
messages = await get_raw_msg_before_timestamp_with_chat(
chat_id=stream_id,
timestamp=time.time(),
limit=20, # 获取更多消息以供筛选
@@ -104,7 +104,7 @@ async def build_cross_context_s4u(
user_name = (
target_user_info.get("person_name") or target_user_info.get("user_nickname") or user_id
)
formatted_messages, _ = build_readable_messages_with_id(
formatted_messages, _ = await build_readable_messages_with_id(
user_messages, timestamp_mode="relative"
)
cross_context_messages.append(
@@ -161,14 +161,14 @@ async def get_chat_history_by_group_name(group_name: str) -> str:
stream_id = found_stream.stream_id
try:
messages = get_raw_msg_before_timestamp_with_chat(
messages = await get_raw_msg_before_timestamp_with_chat(
chat_id=stream_id,
timestamp=time.time(),
limit=5, # 可配置
)
if messages:
chat_name = get_chat_manager().get_stream_name(stream_id) or chat_raw_id
formatted_messages, _ = build_readable_messages_with_id(messages, timestamp_mode="relative")
formatted_messages, _ = await build_readable_messages_with_id(messages, timestamp_mode="relative")
cross_context_messages.append(f'[以下是来自"{chat_name}"的近期消息]\n{formatted_messages}')
except Exception as e:
logger.error(f"获取聊天 {chat_raw_id} 的消息失败: {e}")

46
src/plugin_system/apis/message_api.py
View File

@@ -8,7 +8,7 @@
readable_text = message_api.build_readable_messages(messages)
"""
from typing import List, Dict, Any, Tuple, Optional
from typing import List, Dict, Any, Tuple, Optional, Coroutine
from src.config.config import global_config
import time
from src.chat.utils.chat_message_builder import (
@@ -36,7 +36,7 @@ from src.chat.utils.chat_message_builder import (
def get_messages_by_time(
start_time: float, end_time: float, limit: int = 0, limit_mode: str = "latest", filter_mai: bool = False
) -> List[Dict[str, Any]]:
) -> Coroutine[Any, Any, list[dict[str, Any]]]:
"""
获取指定时间范围内的消息
@@ -62,7 +62,7 @@ def get_messages_by_time(
return get_raw_msg_by_timestamp(start_time, end_time, limit, limit_mode)
def get_messages_by_time_in_chat(
async def get_messages_by_time_in_chat(
chat_id: str,
start_time: float,
end_time: float,
@@ -97,13 +97,13 @@ def get_messages_by_time_in_chat(
if not isinstance(chat_id, str):
raise ValueError("chat_id 必须是字符串类型")
if filter_mai:
return filter_mai_messages(
get_raw_msg_by_timestamp_with_chat(chat_id, start_time, end_time, limit, limit_mode, filter_command)
return await filter_mai_messages(
await get_raw_msg_by_timestamp_with_chat(chat_id, start_time, end_time, limit, limit_mode, filter_command)
)
return get_raw_msg_by_timestamp_with_chat(chat_id, start_time, end_time, limit, limit_mode, filter_command)
return await get_raw_msg_by_timestamp_with_chat(chat_id, start_time, end_time, limit, limit_mode, filter_command)
def get_messages_by_time_in_chat_inclusive(
async def get_messages_by_time_in_chat_inclusive(
chat_id: str,
start_time: float,
end_time: float,
@@ -138,12 +138,12 @@ def get_messages_by_time_in_chat_inclusive(
if not isinstance(chat_id, str):
raise ValueError("chat_id 必须是字符串类型")
if filter_mai:
return filter_mai_messages(
get_raw_msg_by_timestamp_with_chat_inclusive(
return await filter_mai_messages(
await get_raw_msg_by_timestamp_with_chat_inclusive(
chat_id, start_time, end_time, limit, limit_mode, filter_command
)
)
return get_raw_msg_by_timestamp_with_chat_inclusive(
return await get_raw_msg_by_timestamp_with_chat_inclusive(
chat_id, start_time, end_time, limit, limit_mode, filter_command
)
@@ -155,7 +155,7 @@ def get_messages_by_time_in_chat_for_users(
person_ids: List[str],
limit: int = 0,
limit_mode: str = "latest",
) -> List[Dict[str, Any]]:
) -> Coroutine[Any, Any, list[dict[str, Any]]]:
"""
获取指定聊天中指定用户在指定时间范围内的消息
@@ -186,7 +186,7 @@ def get_messages_by_time_in_chat_for_users(
def get_random_chat_messages(
start_time: float, end_time: float, limit: int = 0, limit_mode: str = "latest", filter_mai: bool = False
) -> List[Dict[str, Any]]:
) -> Coroutine[Any, Any, list[dict[str, Any]]]:
"""
随机选择一个聊天，返回该聊天在指定时间范围内的消息
@@ -214,7 +214,7 @@ def get_random_chat_messages(
def get_messages_by_time_for_users(
start_time: float, end_time: float, person_ids: List[str], limit: int = 0, limit_mode: str = "latest"
) -> List[Dict[str, Any]]:
) -> Coroutine[Any, Any, list[dict[str, Any]]]:
"""
获取指定用户在所有聊天中指定时间范围内的消息
@@ -238,7 +238,8 @@ def get_messages_by_time_for_users(
return get_raw_msg_by_timestamp_with_users(start_time, end_time, person_ids, limit, limit_mode)
def get_messages_before_time(timestamp: float, limit: int = 0, filter_mai: bool = False) -> List[Dict[str, Any]]:
def get_messages_before_time(timestamp: float, limit: int = 0, filter_mai: bool = False) -> Coroutine[
Any, Any, list[dict[str, Any]]]:
"""
获取指定时间戳之前的消息
@@ -264,7 +265,7 @@ def get_messages_before_time(timestamp: float, limit: int = 0, filter_mai: bool
def get_messages_before_time_in_chat(
chat_id: str, timestamp: float, limit: int = 0, filter_mai: bool = False
) -> List[Dict[str, Any]]:
) -> Coroutine[Any, Any, list[dict[str, Any]]]:
"""
获取指定聊天中指定时间戳之前的消息
@@ -293,7 +294,8 @@ def get_messages_before_time_in_chat(
return get_raw_msg_before_timestamp_with_chat(chat_id, timestamp, limit)
def get_messages_before_time_for_users(timestamp: float, person_ids: List[str], limit: int = 0) -> List[Dict[str, Any]]:
def get_messages_before_time_for_users(timestamp: float, person_ids: List[str], limit: int = 0) -> Coroutine[
Any, Any, list[dict[str, Any]]]:
"""
获取指定用户在指定时间戳之前的消息
@@ -317,7 +319,7 @@ def get_messages_before_time_for_users(timestamp: float, person_ids: List[str],
def get_recent_messages(
chat_id: str, hours: float = 24.0, limit: int = 100, limit_mode: str = "latest", filter_mai: bool = False
) -> List[Dict[str, Any]]:
) -> Coroutine[Any, Any, list[dict[str, Any]]]:
"""
获取指定聊天中最近一段时间的消息
@@ -354,7 +356,8 @@ def get_recent_messages(
# =============================================================================
def count_new_messages(chat_id: str, start_time: float = 0.0, end_time: Optional[float] = None) -> int:
def count_new_messages(chat_id: str, start_time: float = 0.0, end_time: Optional[float] = None) -> Coroutine[
Any, Any, int]:
"""
计算指定聊天中从开始时间到结束时间的新消息数量
@@ -378,7 +381,8 @@ def count_new_messages(chat_id: str, start_time: float = 0.0, end_time: Optional
return num_new_messages_since(chat_id, start_time, end_time)
def count_new_messages_for_users(chat_id: str, start_time: float, end_time: float, person_ids: List[str]) -> int:
def count_new_messages_for_users(chat_id: str, start_time: float, end_time: float, person_ids: List[str]) -> Coroutine[
Any, Any, int]:
"""
计算指定聊天中指定用户从开始时间到结束时间的新消息数量
@@ -416,7 +420,7 @@ def build_readable_messages_to_str(
read_mark: float = 0.0,
truncate: bool = False,
show_actions: bool = False,
) -> str:
) -> Coroutine[Any, Any, str]:
"""
将消息列表构建成可读的字符串
@@ -478,7 +482,7 @@ async def get_person_ids_from_messages(messages: List[Dict[str, Any]]) -> List[s
# =============================================================================
def filter_mai_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
async def filter_mai_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""
从消息列表中移除麦麦的消息
Args:

356
src/plugin_system/apis/permission_api.py
View File

@@ -1,13 +1,8 @@
"""
权限系统API - 提供权限管理相关的API接口
这个模块提供了权限系统的核心API，包括权限检查、权限节点管理等功能。
插件可以通过这些API来检查用户权限和管理权限节点。
"""
"""纯异步权限API定义。所有外部调用方必须使用 await。"""
from typing import Optional, List, Dict, Any
from enum import Enum
from dataclasses import dataclass
from enum import Enum
from abc import ABC, abstractmethod
from src.common.logger import get_logger
@@ -16,325 +11,172 @@ logger = get_logger(__name__)
class PermissionLevel(Enum):
"""权限等级枚举"""
MASTER = "master" # 最高权限，无视所有权限节点
MASTER = "master"
@dataclass
class PermissionNode:
"""权限节点数据类"""
node_name: str # 权限节点名称，如 "plugin.example.command.test"
description: str # 权限节点描述
plugin_name: str # 所属插件名称
default_granted: bool = False # 默认是否授权
node_name: str
description: str
plugin_name: str
default_granted: bool = False
@dataclass
class UserInfo:
"""用户信息数据类"""
platform: str # 平台类型，如 "qq"
user_id: str # 用户ID
platform: str
user_id: str
def __post_init__(self):
"""确保user_id是字符串类型"""
self.user_id = str(self.user_id)
def to_tuple(self) -> tuple[str, str]:
"""转换为元组格式"""
return (self.platform, self.user_id)
class IPermissionManager(ABC):
"""权限管理器接口"""
@abstractmethod
async def check_permission(self, user: UserInfo, permission_node: str) -> bool: ...
@abstractmethod
def check_permission(self, user: UserInfo, permission_node: str) -> bool:
"""
检查用户是否拥有指定权限节点
Args:
user: 用户信息
permission_node: 权限节点名称
Returns:
bool: 是否拥有权限
"""
pass
def is_master(self, user: UserInfo) -> bool: ... # 同步快速判断
@abstractmethod
def is_master(self, user: UserInfo) -> bool:
"""
检查用户是否为Master用户
Args:
user: 用户信息
Returns:
bool: 是否为Master用户
"""
pass
async def register_permission_node(self, node: PermissionNode) -> bool: ...
@abstractmethod
def register_permission_node(self, node: PermissionNode) -> bool:
"""
注册权限节点
Args:
node: 权限节点
Returns:
bool: 注册是否成功
"""
pass
async def grant_permission(self, user: UserInfo, permission_node: str) -> bool: ...
@abstractmethod
def grant_permission(self, user: UserInfo, permission_node: str) -> bool:
"""
授权用户权限节点
Args:
user: 用户信息
permission_node: 权限节点名称
Returns:
bool: 授权是否成功
"""
pass
async def revoke_permission(self, user: UserInfo, permission_node: str) -> bool: ...
@abstractmethod
def revoke_permission(self, user: UserInfo, permission_node: str) -> bool:
"""
撤销用户权限节点
Args:
user: 用户信息
permission_node: 权限节点名称
Returns:
bool: 撤销是否成功
"""
pass
async def get_user_permissions(self, user: UserInfo) -> List[str]: ...
@abstractmethod
def get_user_permissions(self, user: UserInfo) -> List[str]:
"""
获取用户拥有的所有权限节点
Args:
user: 用户信息
Returns:
List[str]: 权限节点列表
"""
pass
async def get_all_permission_nodes(self) -> List[PermissionNode]: ...
@abstractmethod
def get_all_permission_nodes(self) -> List[PermissionNode]:
"""
获取所有已注册的权限节点
Returns:
List[PermissionNode]: 权限节点列表
"""
pass
@abstractmethod
def get_plugin_permission_nodes(self, plugin_name: str) -> List[PermissionNode]:
"""
获取指定插件的所有权限节点
Args:
plugin_name: 插件名称
Returns:
List[PermissionNode]: 权限节点列表
"""
pass
async def get_plugin_permission_nodes(self, plugin_name: str) -> List[PermissionNode]: ...
class PermissionAPI:
"""权限系统API类"""
def __init__(self):
self._permission_manager: Optional[IPermissionManager] = None
# 需要保留的前缀（视为绝对节点名，不再自动加 plugins.<plugin>. 前缀）
self.RESERVED_PREFIXES: tuple[str, ...] = (
"system.")
# 系统节点列表 (name, description, default_granted)
self._SYSTEM_NODES: list[tuple[str, str, bool]] = [
("system.superuser", "系统超级管理员：拥有所有权限", False),
("system.permission.manage", "系统权限管理：可管理所有权限节点", False),
("system.permission.view", "系统权限查看：可查看所有权限节点", True),
]
self._system_nodes_initialized: bool = False
def set_permission_manager(self, manager: IPermissionManager):
"""设置权限管理器实例"""
self._permission_manager = manager
logger.info("权限管理器已设置")
def _ensure_manager(self):
"""确保权限管理器已设置"""
if self._permission_manager is None:
raise RuntimeError("权限管理器未设置，请先调用 set_permission_manager")
def check_permission(self, platform: str, user_id: str, permission_node: str) -> bool:
"""
检查用户是否拥有指定权限节点
Args:
platform: 平台类型，如 "qq"
user_id: 用户ID
permission_node: 权限节点名称
Returns:
bool: 是否拥有权限
Raises:
RuntimeError: 权限管理器未设置时抛出
"""
async def check_permission(self, platform: str, user_id: str, permission_node: str) -> bool:
self._ensure_manager()
user = UserInfo(platform=platform, user_id=str(user_id))
return self._permission_manager.check_permission(user, permission_node)
return await self._permission_manager.check_permission(UserInfo(platform, user_id), permission_node)
def is_master(self, platform: str, user_id: str) -> bool:
"""
检查用户是否为Master用户
Args:
platform: 平台类型，如 "qq"
user_id: 用户ID
Returns:
bool: 是否为Master用户
Raises:
RuntimeError: 权限管理器未设置时抛出
"""
self._ensure_manager()
user = UserInfo(platform=platform, user_id=str(user_id))
return self._permission_manager.is_master(user)
return self._permission_manager.is_master(UserInfo(platform, user_id))
def register_permission_node(
self, node_name: str, description: str, plugin_name: str, default_granted: bool = False
async def register_permission_node(
self,
node_name: str,
description: str,
plugin_name: str,
default_granted: bool = False,
*,
system: bool = False,
allow_relative: bool = True,
) -> bool:
"""
注册权限节点
Args:
node_name: 权限节点名称，如 "plugin.example.command.test"
description: 权限节点描述
plugin_name: 所属插件名称
default_granted: 默认是否授权
Returns:
bool: 注册是否成功
Raises:
RuntimeError: 权限管理器未设置时抛出
"""
self._ensure_manager()
node = PermissionNode(
node_name=node_name, description=description, plugin_name=plugin_name, default_granted=default_granted
original_name = node_name
if system:
# 系统节点必须以 system./sys./core. 等保留前缀开头
if not node_name.startswith(("system.", "sys.", "core.")):
node_name = f"system.{node_name}" # 自动补 system.
else:
# 普通插件节点：若不以保留前缀开头，并允许相对，则自动加前缀
if allow_relative and not node_name.startswith(self.RESERVED_PREFIXES):
node_name = f"plugins.{plugin_name}.{node_name}"
if original_name != node_name:
logger.debug(f"规范化权限节点 '{original_name}' -> '{node_name}'")
node = PermissionNode(node_name, description, plugin_name, default_granted)
return await self._permission_manager.register_permission_node(node)
async def register_system_permission_node(
self, node_name: str, description: str, default_granted: bool = False
) -> bool:
"""注册系统级权限节点（不绑定具体插件，前缀保持 system./sys./core.）。"""
return await self.register_permission_node(
node_name,
description,
plugin_name="__system__",
default_granted=default_granted,
system=True,
allow_relative=True,
)
return self._permission_manager.register_permission_node(node)
def grant_permission(self, platform: str, user_id: str, permission_node: str) -> bool:
"""
授权用户权限节点
Args:
platform: 平台类型，如 "qq"
user_id: 用户ID
permission_node: 权限节点名称
Returns:
bool: 授权是否成功
Raises:
RuntimeError: 权限管理器未设置时抛出
async def init_system_nodes(self) -> None:
"""初始化默认系统权限节点（幂等）。
在设置 permission_manager 之后且数据库准备好时调用一次即可。
"""
if self._system_nodes_initialized:
return
self._ensure_manager()
user = UserInfo(platform=platform, user_id=str(user_id))
return self._permission_manager.grant_permission(user, permission_node)
for name, desc, granted in self._SYSTEM_NODES:
try:
await self.register_system_permission_node(name, desc, granted)
except Exception as e: # 防御性
logger.warning(f"注册系统权限节点 {name} 失败: {e}")
self._system_nodes_initialized = True
def revoke_permission(self, platform: str, user_id: str, permission_node: str) -> bool:
"""
撤销用户权限节点
Args:
platform: 平台类型，如 "qq"
user_id: 用户ID
permission_node: 权限节点名称
Returns:
bool: 撤销是否成功
Raises:
RuntimeError: 权限管理器未设置时抛出
"""
async def grant_permission(self, platform: str, user_id: str, permission_node: str) -> bool:
self._ensure_manager()
user = UserInfo(platform=platform, user_id=str(user_id))
return self._permission_manager.revoke_permission(user, permission_node)
return await self._permission_manager.grant_permission(UserInfo(platform, user_id), permission_node)
def get_user_permissions(self, platform: str, user_id: str) -> List[str]:
"""
获取用户拥有的所有权限节点
Args:
platform: 平台类型，如 "qq"
user_id: 用户ID
Returns:
List[str]: 权限节点列表
Raises:
RuntimeError: 权限管理器未设置时抛出
"""
async def revoke_permission(self, platform: str, user_id: str, permission_node: str) -> bool:
self._ensure_manager()
user = UserInfo(platform=platform, user_id=str(user_id))
return self._permission_manager.get_user_permissions(user)
return await self._permission_manager.revoke_permission(UserInfo(platform, user_id), permission_node)
def get_all_permission_nodes(self) -> List[Dict[str, Any]]:
"""
获取所有已注册的权限节点
Returns:
List[Dict[str, Any]]: 权限节点列表，每个节点包含 node_name, description, plugin_name, default_granted
Raises:
RuntimeError: 权限管理器未设置时抛出
"""
async def get_user_permissions(self, platform: str, user_id: str) -> List[str]:
self._ensure_manager()
nodes = self._permission_manager.get_all_permission_nodes()
return await self._permission_manager.get_user_permissions(UserInfo(platform, user_id))
async def get_all_permission_nodes(self) -> List[Dict[str, Any]]:
self._ensure_manager()
nodes = await self._permission_manager.get_all_permission_nodes()
return [
{
"node_name": node.node_name,
"description": node.description,
"plugin_name": node.plugin_name,
"default_granted": node.default_granted,
"node_name": n.node_name,
"description": n.description,
"plugin_name": n.plugin_name,
"default_granted": n.default_granted,
}
for node in nodes
for n in nodes
]
def get_plugin_permission_nodes(self, plugin_name: str) -> List[Dict[str, Any]]:
"""
获取指定插件的所有权限节点
Args:
plugin_name: 插件名称
Returns:
List[Dict[str, Any]]: 权限节点列表
Raises:
RuntimeError: 权限管理器未设置时抛出
"""
async def get_plugin_permission_nodes(self, plugin_name: str) -> List[Dict[str, Any]]:
self._ensure_manager()
nodes = self._permission_manager.get_plugin_permission_nodes(plugin_name)
nodes = await self._permission_manager.get_plugin_permission_nodes(plugin_name)
return [
{
"node_name": node.node_name,
"description": node.description,
"plugin_name": node.plugin_name,
"default_granted": node.default_granted,
"node_name": n.node_name,
"description": n.description,
"plugin_name": n.plugin_name,
"default_granted": n.default_granted,
}
for node in nodes
for n in nodes
]
# 全局权限API实例
permission_api = PermissionAPI()

179
src/plugin_system/apis/schedule_api.py Normal file
View File

@@ -0,0 +1,179 @@
"""
日程表与月度计划API模块
专门负责日程和月度计划信息的查询与管理，采用标准Python包设计模式
所有对外接口均为异步函数，以便于插件开发者在异步环境中使用。
使用方式：
import asyncio
from src.plugin_system.apis import schedule_api
async def main():
# 获取今日日程
today_schedule = await schedule_api.get_today_schedule()
if today_schedule:
print("今天的日程:", today_schedule)
# 获取当前活动
current_activity = await schedule_api.get_current_activity()
if current_activity:
print("当前活动:", current_activity)
# 获取本月月度计划
from datetime import datetime
this_month = datetime.now().strftime("%Y-%m")
plans = await schedule_api.get_monthly_plans(this_month)
if plans:
print(f"{this_month} 的月度计划:", [p.plan_text for p in plans])
asyncio.run(main())
"""
from datetime import datetime
from typing import List, Dict, Any, Optional
from src.common.database.sqlalchemy_models import MonthlyPlan
from src.common.logger import get_logger
from src.schedule.database import get_active_plans_for_month
from src.schedule.schedule_manager import schedule_manager
logger = get_logger("schedule_api")
class ScheduleAPI:
"""日程表与月度计划API - 负责日程和计划信息的查询与管理"""
@staticmethod
async def get_today_schedule() -> Optional[List[Dict[str, Any]]]:
"""(异步) 获取今天的日程安排
Returns:
Optional[List[Dict[str, Any]]]: 今天的日程列表，如果未生成或未启用则返回None
"""
try:
logger.debug("[ScheduleAPI] 正在获取今天的日程安排...")
return schedule_manager.today_schedule
except Exception as e:
logger.error(f"[ScheduleAPI] 获取今日日程失败: {e}")
return None
@staticmethod
async def get_current_activity() -> Optional[str]:
"""(异步) 获取当前正在进行的活动
Returns:
Optional[str]: 当前活动名称，如果没有则返回None
"""
try:
logger.debug("[ScheduleAPI] 正在获取当前活动...")
return schedule_manager.get_current_activity()
except Exception as e:
logger.error(f"[ScheduleAPI] 获取当前活动失败: {e}")
return None
@staticmethod
async def regenerate_schedule() -> bool:
"""(异步) 触发后台重新生成今天的日程
Returns:
bool: 是否成功触发
"""
try:
logger.info("[ScheduleAPI] 正在触发后台重新生成日程...")
await schedule_manager.generate_and_save_schedule()
return True
except Exception as e:
logger.error(f"[ScheduleAPI] 触发日程重新生成失败: {e}")
return False
@staticmethod
async def get_monthly_plans(target_month: Optional[str] = None) -> List[MonthlyPlan]:
"""(异步) 获取指定月份的有效月度计划
Args:
target_month (Optional[str]): 目标月份，格式为 "YYYY-MM"。如果为None，则使用当前月份。
Returns:
List[MonthlyPlan]: 月度计划对象列表
"""
if target_month is None:
target_month = datetime.now().strftime("%Y-%m")
try:
logger.debug(f"[ScheduleAPI] 正在获取 {target_month} 的月度计划...")
return await get_active_plans_for_month(target_month)
except Exception as e:
logger.error(f"[ScheduleAPI] 获取 {target_month} 月度计划失败: {e}")
return []
@staticmethod
async def ensure_monthly_plans(target_month: Optional[str] = None) -> bool:
"""(异步) 确保指定月份存在月度计划，如果不存在则触发生成
Args:
target_month (Optional[str]): 目标月份，格式为 "YYYY-MM"。如果为None，则使用当前月份。
Returns:
bool: 操作是否成功 (如果已存在或成功生成)
"""
if target_month is None:
target_month = datetime.now().strftime("%Y-%m")
try:
logger.info(f"[ScheduleAPI] 正在确保 {target_month} 的月度计划存在...")
return await schedule_manager.plan_manager.ensure_and_generate_plans_if_needed(target_month)
except Exception as e:
logger.error(f"[ScheduleAPI] 确保 {target_month} 月度计划失败: {e}")
return False
@staticmethod
async def archive_monthly_plans(target_month: Optional[str] = None) -> bool:
"""(异步) 归档指定月份的月度计划
Args:
target_month (Optional[str]): 目标月份，格式为 "YYYY-MM"。如果为None，则使用当前月份。
Returns:
bool: 操作是否成功
"""
if target_month is None:
target_month = datetime.now().strftime("%Y-%m")
try:
logger.info(f"[ScheduleAPI] 正在归档 {target_month} 的月度计划...")
await schedule_manager.plan_manager.archive_current_month_plans(target_month)
return True
except Exception as e:
logger.error(f"[ScheduleAPI] 归档 {target_month} 月度计划失败: {e}")
return False
# =============================================================================
# 模块级别的便捷函数 (全部为异步)
# =============================================================================
async def get_today_schedule() -> Optional[List[Dict[str, Any]]]:
"""(异步) 获取今天的日程安排的便捷函数"""
return await ScheduleAPI.get_today_schedule()
async def get_current_activity() -> Optional[str]:
"""(异步) 获取当前正在进行的活动的便捷函数"""
return await ScheduleAPI.get_current_activity()
async def regenerate_schedule() -> bool:
"""(异步) 触发后台重新生成今天的日程的便捷函数"""
return await ScheduleAPI.regenerate_schedule()
async def get_monthly_plans(target_month: Optional[str] = None) -> List[MonthlyPlan]:
"""(异步) 获取指定月份的有效月度计划的便捷函数"""
return await ScheduleAPI.get_monthly_plans(target_month)
async def ensure_monthly_plans(target_month: Optional[str] = None) -> bool:
"""(异步) 确保指定月份存在月度计划的便捷函数"""
return await ScheduleAPI.ensure_monthly_plans(target_month)
async def archive_monthly_plans(target_month: Optional[str] = None) -> bool:
"""(异步) 归档指定月份的月度计划的便捷函数"""
return await ScheduleAPI.archive_monthly_plans(target_month)

4
src/plugin_system/apis/send_api.py
View File

@@ -118,10 +118,10 @@ async def wait_adapter_response(request_id: str, timeout: float = 30.0) -> dict:
response = await asyncio.wait_for(future, timeout=timeout)
return response
except asyncio.TimeoutError:
_adapter_response_pool.pop(request_id, None)
await _adapter_response_pool.pop(request_id, None)
return {"status": "error", "message": "timeout"}
except Exception as e:
_adapter_response_pool.pop(request_id, None)
await _adapter_response_pool.pop(request_id, None)
return {"status": "error", "message": str(e)}

6
src/plugin_system/base/base_event.py
View File

@@ -1,5 +1,6 @@
import asyncio
from typing import List, Dict, Any, Optional
from src.common.logger import get_logger
logger = get_logger("base_event")
@@ -90,8 +91,6 @@ class BaseEvent:
self.allowed_subscribers = allowed_subscribers # 记录事件处理器名
self.allowed_triggers = allowed_triggers # 记录插件名
from src.plugin_system.base.base_events_handler import BaseEventHandler
self.subscribers: List["BaseEventHandler"] = [] # 订阅该事件的事件处理器列表
self.event_handle_lock = asyncio.Lock()
@@ -150,7 +149,8 @@ class BaseEvent:
return HandlerResultsCollection(processed_results)
async def _execute_subscriber(self, subscriber, params: dict) -> HandlerResult:
@staticmethod
async def _execute_subscriber(subscriber, params: dict) -> HandlerResult:
"""执行单个订阅者处理器"""
try:
return await subscriber.execute(params)

3
src/plugin_system/base/plugin_base.py
View File

@@ -277,7 +277,8 @@ class PluginBase(ABC):
return config_version_field.default
return "1.0.0"
def _get_current_config_version(self, config: Dict[str, Any]) -> str:
@staticmethod
def _get_current_config_version(config: Dict[str, Any]) -> str:
"""从配置文件中获取当前版本号"""
if "plugin" in config and "config_version" in config["plugin"]:
return str(config["plugin"]["config_version"])

Some files were not shown because too many files have changed in this diff Show More