K2Cr2O1
f38c7cf12a
移除config.py和ai_chat.py中硬编码的DeepSeek API密钥，改为从环境变量获取 简化ai_chat.py的回复逻辑，去除Markdown转换和图片渲染功能
120 lines
4.2 KiB
Python
120 lines
4.2 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
AI 聊天插件，支持向量数据库记忆功能
|
|
"""
|
|
import time
|
|
import uuid
|
|
from core.managers.command_manager import matcher
|
|
from models.events.message import GroupMessageEvent, PrivateMessageEvent
|
|
from core.managers.vectordb_manager import vectordb_manager
|
|
from core.utils.logger import ModuleLogger
|
|
from core.config_loader import global_config
|
|
|
|
logger = ModuleLogger("AIChat")
|
|
|
|
__plugin_meta__ = {
|
|
"name": "AI 聊天",
|
|
"description": "支持向量数据库记忆功能的 AI 聊天助手",
|
|
"usage": "/chat <内容> - 与 AI 进行对话"
|
|
}
|
|
|
|
# 尝试导入 OpenAI 客户端
|
|
try:
|
|
from openai import AsyncOpenAI
|
|
OPENAI_AVAILABLE = True
|
|
except ImportError:
|
|
OPENAI_AVAILABLE = False
|
|
|
|
async def get_ai_response(user_id: int, group_id: int, user_message: str) -> str:
|
|
"""获取 AI 回复，包含向量数据库记忆"""
|
|
if not OPENAI_AVAILABLE:
|
|
return "请先安装 openai 库: pip install openai"
|
|
|
|
# 从配置中获取 DeepSeek API 配置（复用跨平台插件的配置或全局配置）
|
|
api_key = getattr(global_config.cross_platform, 'deepseek_api_key', None) or "sk-f71322a9fbba4b05a7df969cb4004f06"
|
|
api_url = getattr(global_config.cross_platform, 'deepseek_api_url', "https://api.deepseek.com/v1")
|
|
model = getattr(global_config.cross_platform, 'deepseek_model', "deepseek-chat")
|
|
|
|
if api_key == "your-api-key":
|
|
return "请先在配置中设置 DeepSeek API Key"
|
|
|
|
# 1. 从向量数据库检索相关记忆
|
|
collection_name = f"chat_memory_{user_id}"
|
|
memory_context = ""
|
|
|
|
try:
|
|
results = vectordb_manager.query_texts(
|
|
collection_name=collection_name,
|
|
query_texts=[user_message],
|
|
n_results=3
|
|
)
|
|
|
|
if results and results.get("documents") and results["documents"][0]:
|
|
memory_context = "\n\n相关历史记忆：\n"
|
|
for i, doc in enumerate(results["documents"][0], 1):
|
|
memory_context += f"{i}. {doc}\n"
|
|
except Exception as e:
|
|
logger.error(f"检索聊天记忆失败: {e}")
|
|
|
|
# 2. 构建 Prompt
|
|
system_prompt = f"""你是一个友好的 AI 助手。请根据用户的输入进行回复。
|
|
如果提供了相关历史记忆，请参考这些记忆来保持对话的连贯性。{memory_context}"""
|
|
|
|
try:
|
|
client = AsyncOpenAI(
|
|
api_key=api_key,
|
|
base_url=api_url.replace("/chat/completions", "")
|
|
)
|
|
|
|
response = await client.chat.completions.create(
|
|
model=model,
|
|
messages=[
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_message}
|
|
],
|
|
temperature=0.7,
|
|
max_tokens=1000
|
|
)
|
|
|
|
ai_reply = response.choices[0].message.content
|
|
|
|
# 3. 将本次对话存入向量数据库
|
|
if ai_reply:
|
|
try:
|
|
doc_id = str(uuid.uuid4())
|
|
text_to_embed = f"用户: {user_message}\nAI: {ai_reply}"
|
|
metadata = {
|
|
"user_id": user_id,
|
|
"group_id": group_id,
|
|
"timestamp": int(time.time())
|
|
}
|
|
|
|
vectordb_manager.add_texts(
|
|
collection_name=collection_name,
|
|
texts=[text_to_embed],
|
|
metadatas=[metadata],
|
|
ids=[doc_id]
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"保存聊天记忆失败: {e}")
|
|
|
|
return ai_reply
|
|
except Exception as e:
|
|
logger.error(f"AI 聊天请求失败: {e}")
|
|
return f"请求失败: {str(e)}"
|
|
|
|
@matcher.command("chat")
|
|
async def chat_command(event: GroupMessageEvent | PrivateMessageEvent, args: list[str]):
|
|
"""AI 聊天命令"""
|
|
if not args:
|
|
await event.reply("请提供要聊天的内容，例如：/chat 你好")
|
|
return
|
|
|
|
user_message = " ".join(args)
|
|
user_id = event.user_id
|
|
group_id = getattr(event, 'group_id', 0)
|
|
|
|
await event.reply("正在思考中...")
|
|
reply = await get_ai_response(user_id, group_id, user_message)
|
|
await event.reply(reply)
|