NeoBot/plugins/ai_chat.py

# -*- 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")

# 尝试导入 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 "your-api-key"
    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)