This commit is contained in:
SengokuCola
2025-05-05 11:32:21 +08:00
17 changed files with 363 additions and 139 deletions

View File

@@ -174,6 +174,11 @@ def main():
embed_manager.load_from_file()
except Exception as e:
logger.error("从文件加载Embedding库时发生错误{}".format(e))
if "嵌入模型与本地存储不一致" in str(e):
logger.error("检测到嵌入模型与本地存储不一致,已终止导入。请检查模型设置或清空嵌入库后重试。")
logger.error("请保证你的嵌入模型从未更改,并且在导入时使用相同的模型")
# print("检测到嵌入模型与本地存储不一致,已终止导入。请检查模型设置或清空嵌入库后重试。")
sys.exit(1)
logger.error("如果你是第一次导入知识,请忽略此错误")
logger.info("Embedding库加载完成")
# 初始化KG

View File

@@ -1,8 +0,0 @@
from fastapi import FastAPI
from strawberry.fastapi import GraphQLRouter
app = FastAPI()
graphql_router = GraphQLRouter(schema=None, path="/") # Replace `None` with your actual schema
app.include_router(graphql_router, prefix="/graphql", tags=["GraphQL"])

16
src/api/apiforgui.py Normal file
View File

@@ -0,0 +1,16 @@
from src.heart_flow.heartflow import heartflow
from src.heart_flow.sub_heartflow import ChatState
async def get_all_subheartflow_ids() -> list:
"""获取所有子心流的ID列表"""
all_subheartflows = heartflow.subheartflow_manager.get_all_subheartflows()
return [subheartflow.subheartflow_id for subheartflow in all_subheartflows]
async def forced_change_subheartflow_status(subheartflow_id: str, status: ChatState) -> bool:
"""强制改变子心流的状态"""
subheartflow = await heartflow.get_or_create_subheartflow(subheartflow_id)
if subheartflow:
return await heartflow.force_change_subheartflow_status(subheartflow_id, status)
return False

View File

@@ -1,155 +1,187 @@
from typing import List, Optional
from typing import List, Optional, Dict, Any
import strawberry
# from packaging.version import Version, InvalidVersion
# from packaging.specifiers import SpecifierSet, InvalidSpecifier
# from ..config.config import global_config
# import os
from packaging.version import Version
import os
ROOT_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
@strawberry.type
class BotConfig:
class APIBotConfig:
"""机器人配置类"""
INNER_VERSION: Version
INNER_VERSION: Version # 配置文件内部版本号
MAI_VERSION: str # 硬编码的版本信息
# bot
BOT_QQ: Optional[int]
BOT_NICKNAME: Optional[str]
BOT_ALIAS_NAMES: List[str] # 别名,可以通过这个叫它
BOT_QQ: Optional[int] # 机器人QQ号
BOT_NICKNAME: Optional[str] # 机器人昵称
BOT_ALIAS_NAMES: List[str] # 机器人别名列表
# group
talk_allowed_groups: set
talk_frequency_down_groups: set
ban_user_id: set
talk_allowed_groups: List[int] # 允许回复消息的群号列表
talk_frequency_down_groups: List[int] # 降低回复频率的群号列表
ban_user_id: List[int] # 禁止回复和读取消息的QQ号列表
# personality
personality_core: str # 建议20字以内谁再写3000字小作文敲谁脑袋
personality_sides: List[str]
personality_core: str # 人格核心特点描述
personality_sides: List[str] # 人格细节描述列表
# identity
identity_detail: List[str]
height: int # 身高 单位厘米
weight: int # 体重 单位千克
age: int # 年龄 单位岁
identity_detail: List[str] # 身份特点列表
age: int # 年龄(岁)
gender: str # 性别
appearance: str # 外貌特征
appearance: str # 外貌特征描述
# schedule
ENABLE_SCHEDULE_GEN: bool # 是否启用日程生成
PROMPT_SCHEDULE_GEN: str
SCHEDULE_DOING_UPDATE_INTERVAL: int # 日程表更新间隔 单位秒
SCHEDULE_TEMPERATURE: float # 日程表温度建议0.5-1.0
ENABLE_SCHEDULE_INTERACTION: bool # 是否启用日程交互
PROMPT_SCHEDULE_GEN: str # 日程生成提示词
SCHEDULE_DOING_UPDATE_INTERVAL: int # 日程进行中更新间隔
SCHEDULE_TEMPERATURE: float # 日程生成温度
TIME_ZONE: str # 时区
# message
MAX_CONTEXT_SIZE: int # 上下文最大消息数
emoji_chance: float # 发送表情包的基础概率
thinking_timeout: int # 思考时间
model_max_output_length: int # 最大回复长度
message_buffer: bool # 消息缓冲器
# platforms
platforms: Dict[str, str] # 平台信息
ban_words: set
ban_msgs_regex: set
# heartflow
# enable_heartflow: bool = False # 是否启用心流
sub_heart_flow_update_interval: int # 子心流更新频率,间隔 单位秒
sub_heart_flow_freeze_time: int # 子心流冻结时间,超过这个时间没有回复,子心流会冻结,间隔 单位秒
sub_heart_flow_stop_time: int # 子心流停止时间,超过这个时间没有回复,子心流会停止,间隔 单位秒
heart_flow_update_interval: int # 心流更新频率,间隔 单位秒
observation_context_size: int # 心流观察到的最长上下文大小,超过这个值的上下文会被压缩
compressed_length: int # 不能大于observation_context_size,心流上下文压缩的最短压缩长度超过心流观察到的上下文长度会压缩最短压缩长度为5
compress_length_limit: int # 最多压缩份数,超过该数值的压缩上下文会被删除
# chat
allow_focus_mode: bool # 是否允许专注模式
base_normal_chat_num: int # 基础普通聊天次数
base_focused_chat_num: int # 基础专注聊天次数
observation_context_size: int # 观察上下文大小
message_buffer: bool # 是否启用消息缓冲
ban_words: List[str] # 禁止词列表
ban_msgs_regex: List[str] # 禁止消息的正则表达式列表
# willing
# normal_chat
MODEL_R1_PROBABILITY: float # 模型推理概率
MODEL_V3_PROBABILITY: float # 模型普通概率
emoji_chance: float # 表情符号出现概率
thinking_timeout: int # 思考超时时间
willing_mode: str # 意愿模式
response_willing_amplifier: float # 回复意愿放大系数
response_interested_rate_amplifier: float # 回复兴趣放大系数
down_frequency_rate: float # 降低回复频率的群组回复意愿降低系数
emoji_response_penalty: float # 表情回复惩罚
mentioned_bot_inevitable_reply: bool # 提及 bot 必然回复
at_bot_inevitable_reply: bool # @bot 必然回复
response_willing_amplifier: float # 回复意愿放大
response_interested_rate_amplifier: float # 回复兴趣放大
down_frequency_rate: float # 降低频率率
emoji_response_penalty: float # 表情回复惩罚
mentioned_bot_inevitable_reply: bool # 提到机器人时是否必定回复
at_bot_inevitable_reply: bool # @机器人时是否必定回复
# response
response_mode: str # 回复策略
MODEL_R1_PROBABILITY: float # R1模型概
MODEL_V3_PROBABILITY: float # V3模型概率
# MODEL_R1_DISTILL_PROBABILITY: float # R1蒸馏模型概率
# focus_chat
reply_trigger_threshold: float # 回复触发阈值
default_decay_rate_per_second: float # 默认每秒衰减
consecutive_no_reply_threshold: int # 连续不回复阈值
# compressed
compressed_length: int # 压缩长度
compress_length_limit: int # 压缩长度限制
# emoji
max_emoji_num: int # 表情包最大数量
max_reach_deletion: bool # 开启则在达到最大数量时删除表情包,关闭则不会继续收集表情包
EMOJI_CHECK_INTERVAL: int # 表情检查间隔(分钟)
EMOJI_REGISTER_INTERVAL: int # 表情注册间隔(分钟
EMOJI_SAVE: bool # 表情
EMOJI_CHECK: bool # 是否开启过滤
EMOJI_CHECK_PROMPT: str # 表情包过滤要求
max_emoji_num: int # 最大表情符号数量
max_reach_deletion: bool # 达到最大数量时是否删除
EMOJI_CHECK_INTERVAL: int # 表情检查间隔
EMOJI_REGISTER_INTERVAL: Optional[int] # 表情注册间隔(兼容性保留
EMOJI_SAVE: bool # 是否保存表情
EMOJI_CHECK: bool # 是否检查表情
EMOJI_CHECK_PROMPT: str # 表情检查提示词
# memory
build_memory_interval: int # 记忆构建间隔(秒)
memory_build_distribution: list # 记忆构建分布参数分布1均值标准差权重分布2均值标准差权重
build_memory_sample_num: int # 记忆构建采样数量
build_memory_sample_length: int # 记忆构建采样长度
build_memory_interval: int # 构建记忆间隔
memory_build_distribution: List[float] # 记忆构建分布
build_memory_sample_num: int # 构建记忆样本数量
build_memory_sample_length: int # 构建记忆样本长度
memory_compress_rate: float # 记忆压缩率
forget_memory_interval: int # 记忆遗忘间隔(秒)
memory_forget_time: int # 记忆遗忘时间(小时)
memory_forget_percentage: float # 记忆遗忘比例
memory_ban_words: list # 添加新的配置项默认值
forget_memory_interval: int # 忘记记忆间隔
memory_forget_time: int # 记忆忘记时间
memory_forget_percentage: float # 记忆忘记百分比
consolidate_memory_interval: int # 巩固记忆间隔
consolidation_similarity_threshold: float # 巩固相似度阈值
consolidation_check_percentage: float # 巩固检查百分比
memory_ban_words: List[str] # 记忆禁止词列表
# mood
mood_update_interval: float # 情绪更新间隔 单位秒
mood_update_interval: float # 情绪更新间隔
mood_decay_rate: float # 情绪衰减率
mood_intensity_factor: float # 情绪强度因子
# keywords
keywords_reaction_rules: list # 关键词回复规则
# keywords_reaction
keywords_reaction_enable: bool # 是否启用关键词反应
keywords_reaction_rules: List[Dict[str, Any]] # 关键词反应规则
# chinese_typo
chinese_typo_enable: bool # 是否启用中文错别字生成器
chinese_typo_error_rate: float # 单字替换概
chinese_typo_min_freq: int # 最小字频阈值
chinese_typo_tone_error_rate: float # 声调错误
chinese_typo_word_replace_rate: float # 词替换
chinese_typo_enable: bool # 是否启用中文错别字
chinese_typo_error_rate: float # 中文错别字错误
chinese_typo_min_freq: int # 中文错别字最小频率
chinese_typo_tone_error_rate: float # 中文错别字声调错误率
chinese_typo_word_replace_rate: float # 中文错别字单词替换率
# response_splitter
enable_response_splitter: bool # 是否启用回复分割器
response_max_length: int # 回复允许的最大长度
response_max_sentence_num: int # 回复允许的最大句子数
response_max_length: int # 回复最大长度
response_max_sentence_num: int # 回复最大句子数
enable_kaomoji_protection: bool # 是否启用颜文字保护
model_max_output_length: int # 模型最大输出长度
# remote
remote_enable: bool # 是否启用远程控制
remote_enable: bool # 是否启用远程功能
# experimental
enable_friend_chat: bool # 是否启用好友聊天
# enable_think_flow: bool # 是否启用思考流程
talk_allowed_private: List[int] # 允许私聊的QQ号列表
enable_pfc_chatting: bool # 是否启用PFC聊天
# 模型配置
llm_reasoning: dict[str, str] # LLM推理
# llm_reasoning_minor: dict[str, str]
llm_normal: dict[str, str] # LLM普通
llm_topic_judge: dict[str, str] # LLM话题判断
llm_summary: dict[str, str] # LLM话题总结
llm_emotion_judge: dict[str, str] # LLM情感判断
embedding: dict[str, str] # 嵌入
vlm: dict[str, str] # VLM
moderation: dict[str, str] # 审核
llm_reasoning: Dict[str, Any] # 推理模型配置
llm_normal: Dict[str, Any] # 普通模型配置
llm_topic_judge: Dict[str, Any] # 主题判断模型配置
llm_summary: Dict[str, Any] # 总结模型配置
llm_emotion_judge: Optional[Dict[str, Any]] # 情绪判断模型配置(兼容性保留)
embedding: Dict[str, Any] # 嵌入模型配置
vlm: Dict[str, Any] # VLM模型配置
moderation: Optional[Dict[str, Any]] # 审核模型配置(兼容性保留)
llm_observation: Dict[str, Any] # 观察模型配置
llm_sub_heartflow: Dict[str, Any] # 子心流模型配置
llm_heartflow: Dict[str, Any] # 心流模型配置
llm_plan: Optional[Dict[str, Any]] # 计划模型配置
llm_PFC_action_planner: Optional[Dict[str, Any]] # PFC行动计划模型配置
llm_PFC_chat: Optional[Dict[str, Any]] # PFC聊天模型配置
llm_PFC_reply_checker: Optional[Dict[str, Any]] # PFC回复检查模型配置
llm_tool_use: Optional[Dict[str, Any]] # 工具使用模型配置
# 实验性
llm_observation: dict[str, str] # LLM观察
llm_sub_heartflow: dict[str, str] # LLM子心流
llm_heartflow: dict[str, str] # LLM心流
api_urls: dict[str, str] # API URLs
api_urls: Optional[Dict[str, str]] # API地址配置
@strawberry.type
class EnvConfig:
pass
class APIEnvConfig:
"""环境变量配置"""
HOST: str # 服务主机地址
PORT: int # 服务端口
PLUGINS: List[str] # 插件列表
MONGODB_HOST: str # MongoDB 主机地址
MONGODB_PORT: int # MongoDB 端口
DATABASE_NAME: str # 数据库名称
CHAT_ANY_WHERE_BASE_URL: str # ChatAnywhere 基础URL
SILICONFLOW_BASE_URL: str # SiliconFlow 基础URL
DEEP_SEEK_BASE_URL: str # DeepSeek 基础URL
DEEP_SEEK_KEY: Optional[str] # DeepSeek API Key
CHAT_ANY_WHERE_KEY: Optional[str] # ChatAnywhere API Key
SILICONFLOW_KEY: Optional[str] # SiliconFlow API Key
SIMPLE_OUTPUT: Optional[bool] # 是否简化输出
CONSOLE_LOG_LEVEL: Optional[str] # 控制台日志等级
FILE_LOG_LEVEL: Optional[str] # 文件日志等级
DEFAULT_CONSOLE_LOG_LEVEL: Optional[str] # 默认控制台日志等级
DEFAULT_FILE_LOG_LEVEL: Optional[str] # 默认文件日志等级
@strawberry.field
def get_env(self) -> str:
return "env"
print("当前路径:")
print(ROOT_PATH)

56
src/api/main.py Normal file
View File

@@ -0,0 +1,56 @@
from fastapi import APIRouter
from strawberry.fastapi import GraphQLRouter
# from src.config.config import BotConfig
from src.common.logger_manager import get_logger
from src.api.reload_config import reload_config as reload_config_func
from src.common.server import global_server
from .apiforgui import get_all_subheartflow_ids, forced_change_subheartflow_status
from src.heart_flow.sub_heartflow import ChatState
# import uvicorn
# import os
router = APIRouter()
logger = get_logger("api")
# maiapi = FastAPI()
logger.info("麦麦API服务器已启动")
graphql_router = GraphQLRouter(schema=None, path="/") # Replace `None` with your actual schema
router.include_router(graphql_router, prefix="/graphql", tags=["GraphQL"])
@router.post("/config/reload")
async def reload_config():
return await reload_config_func()
@router.get("/gui/subheartflow/get/all")
async def get_subheartflow_ids():
"""获取所有子心流的ID列表"""
return await get_all_subheartflow_ids()
@router.post("/gui/subheartflow/forced_change_status")
async def forced_change_subheartflow_status_api(subheartflow_id: str, status: ChatState): # noqa
"""强制改变子心流的状态"""
# 参数检查
if not isinstance(status, ChatState):
logger.warning(f"无效的状态参数: {status}")
return {"status": "failed", "reason": "invalid status"}
logger.info(f"尝试将子心流 {subheartflow_id} 状态更改为 {status.value}")
success = await forced_change_subheartflow_status(subheartflow_id, status)
if success:
logger.info(f"子心流 {subheartflow_id} 状态更改为 {status.value} 成功")
return {"status": "success"}
else:
logger.error(f"子心流 {subheartflow_id} 状态更改为 {status.value} 失败")
return {"status": "failed"}
def start_api_server():
"""启动API服务器"""
global_server.register_router(router, prefix="/api/v1")

24
src/api/reload_config.py Normal file
View File

@@ -0,0 +1,24 @@
from fastapi import HTTPException
from rich.traceback import install
from src.config.config import BotConfig
from src.common.logger_manager import get_logger
import os
install(extra_lines=3)
logger = get_logger("api")
async def reload_config():
try:
from src.config import config as config_module
logger.debug("正在重载配置文件...")
bot_config_path = os.path.join(BotConfig.get_config_dir(), "bot_config.toml")
config_module.global_config = BotConfig.load_config(config_path=bot_config_path)
logger.debug("配置文件重载成功")
return {"status": "reloaded"}
except FileNotFoundError as e:
raise HTTPException(status_code=404, detail=str(e)) from e
except Exception as e:
raise HTTPException(status_code=500, detail=f"重载配置时发生错误: {str(e)}") from e

View File

@@ -825,6 +825,22 @@ INIT_STYLE_CONFIG = {
},
}
API_SERVER_STYLE_CONFIG = {
"advanced": {
"console_format": (
"<white>{time:YYYY-MM-DD HH:mm:ss}</white> | "
"<level>{level: <8}</level> | "
"<light-yellow>API服务</light-yellow> | "
"<level>{message}</level>"
),
"file_format": "{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {extra[module]: <15} | API服务 | {message}",
},
"simple": {
"console_format": "<level>{time:MM-DD HH:mm}</level> | <light-green>API服务</light-green> | {message}",
"file_format": "{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {extra[module]: <15} | API服务 | {message}",
},
}
# 根据SIMPLE_OUTPUT选择配置
MAIN_STYLE_CONFIG = MAIN_STYLE_CONFIG["simple"] if SIMPLE_OUTPUT else MAIN_STYLE_CONFIG["advanced"]
@@ -895,11 +911,11 @@ CHAT_MESSAGE_STYLE_CONFIG = (
)
CHAT_IMAGE_STYLE_CONFIG = CHAT_IMAGE_STYLE_CONFIG["simple"] if SIMPLE_OUTPUT else CHAT_IMAGE_STYLE_CONFIG["advanced"]
INIT_STYLE_CONFIG = INIT_STYLE_CONFIG["simple"] if SIMPLE_OUTPUT else INIT_STYLE_CONFIG["advanced"]
API_SERVER_STYLE_CONFIG = API_SERVER_STYLE_CONFIG["simple"] if SIMPLE_OUTPUT else API_SERVER_STYLE_CONFIG["advanced"]
INTEREST_CHAT_STYLE_CONFIG = (
INTEREST_CHAT_STYLE_CONFIG["simple"] if SIMPLE_OUTPUT else INTEREST_CHAT_STYLE_CONFIG["advanced"]
)
def is_registered_module(record: dict) -> bool:
"""检查是否为已注册的模块"""
return record["extra"].get("module") in _handler_registry

View File

@@ -42,6 +42,7 @@ from src.common.logger import (
CHAT_IMAGE_STYLE_CONFIG,
INIT_STYLE_CONFIG,
INTEREST_CHAT_STYLE_CONFIG,
API_SERVER_STYLE_CONFIG,
)
# 可根据实际需要补充更多模块配置
@@ -88,6 +89,7 @@ MODULE_LOGGER_CONFIGS = {
"chat_image": CHAT_IMAGE_STYLE_CONFIG, # 聊天图片
"init": INIT_STYLE_CONFIG, # 初始化
"interest_chat": INTEREST_CHAT_STYLE_CONFIG, # 兴趣
"api": API_SERVER_STYLE_CONFIG, # API服务器
# ...如有更多模块,继续添加...
}

View File

@@ -1,4 +1,4 @@
from src.heart_flow.sub_heartflow import SubHeartflow
from src.heart_flow.sub_heartflow import SubHeartflow, ChatState
from src.plugins.models.utils_model import LLMRequest
from src.config.config import global_config
from src.plugins.schedule.schedule_generator import bot_schedule
@@ -62,6 +62,11 @@ class Heartflow:
# 不再需要传入 self.current_state
return await self.subheartflow_manager.get_or_create_subheartflow(subheartflow_id)
async def force_change_subheartflow_status(self, subheartflow_id: str, status: ChatState) -> None:
"""强制改变子心流的状态"""
# 这里的 message 是可选的,可能是一个消息对象,也可能是其他类型的数据
return await self.subheartflow_manager.force_change_state(subheartflow_id, status)
async def heartflow_start_working(self):
"""启动后台任务"""
await self.background_task_manager.start_tasks()

View File

@@ -83,6 +83,17 @@ class SubHeartflowManager:
request_type="subheartflow_state_eval", # 保留特定的请求类型
)
async def force_change_state(self, subflow_id: Any, target_state: ChatState) -> bool:
"""强制改变指定子心流的状态"""
async with self._lock:
subflow = self.subheartflows.get(subflow_id)
if not subflow:
logger.warning(f"[强制状态转换]尝试转换不存在的子心流{subflow_id}{target_state.value}")
return False
await subflow.change_chat_state(target_state)
logger.info(f"[强制状态转换]子心流 {subflow_id} 已转换到 {target_state.value}")
return True
def get_all_subheartflows(self) -> List["SubHeartflow"]:
"""获取所有当前管理的 SubHeartflow 实例列表 (快照)。"""
return list(self.subheartflows.values())
@@ -92,7 +103,7 @@ class SubHeartflowManager:
Args:
subheartflow_id: 子心流唯一标识符
# mai_states 参数已被移除,使用 self.mai_state_info
mai_states 参数已被移除,使用 self.mai_state_info
Returns:
成功返回SubHeartflow实例失败返回None
@@ -165,7 +176,7 @@ class SubHeartflowManager:
def get_inactive_subheartflows(self, max_age_seconds=INACTIVE_THRESHOLD_SECONDS):
"""识别并返回需要清理的不活跃(处于ABSENT状态超过一小时)子心流(id, 原因)"""
current_time = time.time()
_current_time = time.time()
flows_to_stop = []
for subheartflow_id, subheartflow in list(self.subheartflows.items()):
@@ -173,8 +184,7 @@ class SubHeartflowManager:
if state != ChatState.ABSENT:
continue
subheartflow.update_last_chat_state_time()
absent_last_time = subheartflow.chat_state_last_time
if max_age_seconds and (current_time - absent_last_time) > max_age_seconds:
_absent_last_time = subheartflow.chat_state_last_time
flows_to_stop.append(subheartflow_id)
return flows_to_stop

View File

@@ -18,6 +18,7 @@ from .plugins.remote import heartbeat_thread # noqa: F401
from .individuality.individuality import Individuality
from .common.server import global_server
from rich.traceback import install
from .api.main import start_api_server
install(extra_lines=3)
@@ -54,6 +55,9 @@ class MainSystem:
self.llm_stats.start()
logger.success("LLM统计功能启动成功")
# 启动API服务器
start_api_server()
logger.success("API服务器启动成功")
# 初始化表情管理器
emoji_manager.initialize()
logger.success("表情包管理器初始化成功")

View File

@@ -1 +0,0 @@

View File

@@ -1,19 +0,0 @@
from fastapi import APIRouter, HTTPException
from rich.traceback import install
install(extra_lines=3)
# 创建APIRouter而不是FastAPI实例
router = APIRouter()
@router.post("/reload-config")
async def reload_config():
try: # TODO: 实现配置重载
# bot_config_path = os.path.join(BotConfig.get_config_dir(), "bot_config.toml")
# BotConfig.reload_config(config_path=bot_config_path)
return {"message": "TODO: 实现配置重载", "status": "unimplemented"}
except FileNotFoundError as e:
raise HTTPException(status_code=404, detail=str(e)) from e
except Exception as e:
raise HTTPException(status_code=500, detail=f"重载配置时发生错误: {str(e)}") from e

View File

@@ -1,4 +0,0 @@
import requests
response = requests.post("http://localhost:8080/api/reload-config")
print(response.json())

View File

@@ -1,6 +1,7 @@
from dataclasses import dataclass
import json
import os
import math
from typing import Dict, List, Tuple
import numpy as np
@@ -25,9 +26,39 @@ from rich.progress import (
)
install(extra_lines=3)
ROOT_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", ".."))
TOTAL_EMBEDDING_TIMES = 3 # 统计嵌入次数
# 嵌入模型测试字符串,测试模型一致性,来自开发群的聊天记录
# 这些字符串的嵌入结果应该是固定的,不能随时间变化
EMBEDDING_TEST_STRINGS = [
"阿卡伊真的太好玩了,神秘性感大女同等着你",
"你怎么知道我arc12.64了",
"我是蕾缪乐小姐的狗",
"关注Oct谢谢喵",
"不是w6我不草",
"关注千石可乐谢谢喵",
"来玩CLANNADAIR樱之诗樱之刻谢谢喵",
"关注墨梓柒谢谢喵",
"Ciallo~",
"来玩巧克甜恋谢谢喵",
"水印",
"我也在纠结晚饭,铁锅炒鸡听着就香!",
"test你妈喵",
]
EMBEDDING_TEST_FILE = os.path.join(ROOT_PATH, "data", "embedding_model_test.json")
EMBEDDING_SIM_THRESHOLD = 0.99
def cosine_similarity(a, b):
# 计算余弦相似度
dot = sum(x * y for x, y in zip(a, b))
norm_a = math.sqrt(sum(x * x for x in a))
norm_b = math.sqrt(sum(x * x for x in b))
if norm_a == 0 or norm_b == 0:
return 0.0
return dot / (norm_a * norm_b)
@dataclass
class EmbeddingStoreItem:
@@ -64,6 +95,46 @@ class EmbeddingStore:
def _get_embedding(self, s: str) -> List[float]:
return self.llm_client.send_embedding_request(global_config["embedding"]["model"], s)
def get_test_file_path(self):
return EMBEDDING_TEST_FILE
def save_embedding_test_vectors(self):
"""保存测试字符串的嵌入到本地"""
test_vectors = {}
for idx, s in enumerate(EMBEDDING_TEST_STRINGS):
test_vectors[str(idx)] = self._get_embedding(s)
with open(self.get_test_file_path(), "w", encoding="utf-8") as f:
json.dump(test_vectors, f, ensure_ascii=False, indent=2)
def load_embedding_test_vectors(self):
"""加载本地保存的测试字符串嵌入"""
path = self.get_test_file_path()
if not os.path.exists(path):
return None
with open(path, "r", encoding="utf-8") as f:
return json.load(f)
def check_embedding_model_consistency(self):
"""校验当前模型与本地嵌入模型是否一致"""
local_vectors = self.load_embedding_test_vectors()
if local_vectors is None:
logger.warning("未检测到本地嵌入模型测试文件,将保存当前模型的测试嵌入。")
self.save_embedding_test_vectors()
return True
for idx, s in enumerate(EMBEDDING_TEST_STRINGS):
local_emb = local_vectors.get(str(idx))
if local_emb is None:
logger.warning("本地嵌入模型测试文件缺失部分测试字符串,将重新保存。")
self.save_embedding_test_vectors()
return True
new_emb = self._get_embedding(s)
sim = cosine_similarity(local_emb, new_emb)
if sim < EMBEDDING_SIM_THRESHOLD:
logger.error("嵌入模型一致性校验失败")
return False
logger.info("嵌入模型一致性校验通过。")
return True
def batch_insert_strs(self, strs: List[str], times: int) -> None:
"""向库中存入字符串"""
total = len(strs)
@@ -216,6 +287,17 @@ class EmbeddingManager:
)
self.stored_pg_hashes = set()
def check_all_embedding_model_consistency(self):
"""对所有嵌入库做模型一致性校验"""
for store in [
self.paragraphs_embedding_store,
self.entities_embedding_store,
self.relation_embedding_store,
]:
if not store.check_embedding_model_consistency():
return False
return True
def _store_pg_into_embedding(self, raw_paragraphs: Dict[str, str]):
"""将段落编码存入Embedding库"""
self.paragraphs_embedding_store.batch_insert_strs(list(raw_paragraphs.values()), times=1)
@@ -239,6 +321,8 @@ class EmbeddingManager:
def load_from_file(self):
"""从文件加载"""
if not self.check_all_embedding_model_consistency():
raise Exception("嵌入模型与本地存储不一致,请检查模型设置或清空嵌入库后重试。")
self.paragraphs_embedding_store.load_from_file()
self.entities_embedding_store.load_from_file()
self.relation_embedding_store.load_from_file()
@@ -250,6 +334,8 @@ class EmbeddingManager:
raw_paragraphs: Dict[str, str],
triple_list_data: Dict[str, List[List[str]]],
):
if not self.check_all_embedding_model_consistency():
raise Exception("嵌入模型与本地存储不一致,请检查模型设置或清空嵌入库后重试。")
"""存储新的数据集"""
self._store_pg_into_embedding(raw_paragraphs)
self._store_ent_into_embedding(triple_list_data)