ai-chat-system/services/matrix_service.py

"""
Matrix Bot 服务 - 处理Matrix消息收发
所有消息都同步到主用户的最新会话
"""
import asyncio
import json
from typing import Optional, Callable
from nio import AsyncClient, MatrixRoom, RoomMessageText, LoginResponse
import logging

from models import SessionLocal, User, Conversation, Message, SystemConfig
from services.conversation_service import ConversationService

logger = logging.getLogger(__name__)

# 主用户ID
MAIN_USER_ID = "main_user"


class MatrixBot:
    def __init__(self):
        self.client: Optional[AsyncClient] = None
        self.homeserver: str = ""
        self.username: str = ""
        self.password: str = ""
        self.access_token: str = ""
        self.is_running: bool = False
        self.on_message_callback: Optional[Callable] = None
    
    async def init_from_config(self):
        """从数据库配置初始化"""
        db = SessionLocal()
        try:
            configs = {c.key: c.value for c in db.query(SystemConfig).all()}
            self.homeserver = configs.get('matrix_homeserver', 'https://matrix.tphai.com')
            self.username = configs.get('matrix_username', '')
            self.password = configs.get('matrix_password', '')
            self.access_token = configs.get('matrix_access_token', '')
            
            if self.username and (self.password or self.access_token):
                await self.connect()
        finally:
            db.close()
    
    async def connect(self):
        """连接到Matrix服务器"""
        if not self.homeserver or not self.username:
            logger.warning("Matrix配置不完整，跳过连接")
            return False
        
        try:
            self.client = AsyncClient(
                self.homeserver, 
                self.username,
                store_path="/tmp/matrix_store"
            )
            
            if self.access_token:
                self.client.access_token = self.access_token
                self.client.user_id = self.username
                logger.info(f"Matrix已设置access_token: {self.username}")
                self.is_running = True
                return True
            
            response = await self.client.login(self.password)
            
            if isinstance(response, LoginResponse):
                logger.info(f"Matrix连接成功: {self.username}")
                self.is_running = True
                return True
            else:
                logger.error(f"Matrix登录失败: {response}")
                return False
        except Exception as e:
            logger.error(f"Matrix连接错误: {e}")
            return False
    
    async def start_sync(self, message_handler: Callable = None):
        """开始同步消息"""
        if not self.client:
            logger.warning("Matrix客户端未初始化")
            return
        
        self.on_message_callback = message_handler
        
        # 注册消息处理器
        self.client.add_event_callback(self._handle_room_message, RoomMessageText)
        
        try:
            await self.client.sync(timeout=10000)
            logger.info("Matrix初始同步完成")
        except Exception as e:
            logger.warning(f"Matrix初始同步失败: {e}")
        
        asyncio.create_task(self._sync_loop())
        logger.info("Matrix同步任务已启动")
    
    async def _sync_loop(self):
        """后台同步循环"""
        while self.is_running:
            try:
                await self.client.sync(timeout=5000)
                await asyncio.sleep(1)
            except Exception as e:
                logger.warning(f"Matrix同步错误: {e}")
                await asyncio.sleep(5)
    
    async def _handle_room_message(self, room: MatrixRoom, event: RoomMessageText):
        """处理Matrix消息 - 所有消息同步到主用户最新会话"""
        # 忽略自己发送的消息
        if event.sender == self.username:
            return
        
        message_text = event.body.strip()
        
        db = SessionLocal()
        try:
            conv_service = ConversationService(db)
            
            # 使用固定主用户
            main_user = conv_service.get_or_create_user(
                MAIN_USER_ID, 
                display_name="主用户", 
                user_type='web'
            )
            
            # 处理 /new 命令 - 创建新会话，不回复
            if message_text == "/new":
                conversation = conv_service.create_conversation(main_user.id)
                await self.send_message(room.room_id, "✅ 已创建新会话")
                logger.info(f"Matrix创建新会话: {conversation.conversation_id}")
                
                # 通知WebSocket客户端
                if self.on_message_callback:
                    await self.on_message_callback(
                        action="new_conversation",
                        conversation_id=conversation.conversation_id,
                        room_id=room.room_id
                    )
                return
            
            # 获取最新会话
            conversations = conv_service.get_user_conversations(main_user.id)
            if not conversations:
                # 如果没有会话，创建一个
                conversation = conv_service.create_conversation(main_user.id)
            else:
                conversation = conversations[0]  # 最新会话
            
            # 保存用户消息
            user_msg = conv_service.add_message(
                conversation_id=conversation.id,
                role='user',
                content=message_text,
                source='matrix',
                extra_data={
                    'event_id': event.event_id, 
                    'room_id': room.room_id,
                    'sender': event.sender
                }
            )
            
            logger.info(f"Matrix消息保存到会话 {conversation.conversation_id}")
            
            # 调用消息处理器（获取AI回复并同步）
            if self.on_message_callback:
                await self.on_message_callback(
                    action="chat",
                    conversation_id=conversation.conversation_id,
                    user_message=message_text,
                    room_id=room.room_id,
                    message_id=user_msg.id
                )
        finally:
            db.close()
    
    async def send_message(self, room_id: str, message: str):
        """发送消息到Matrix房间"""
        if not self.client:
            return False
        
        try:
            await self.client.room_send(
                room_id=room_id,
                message_type="m.room.message",
                content={
                    "msgtype": "m.text",
                    "body": message
                }
            )
            return True
        except Exception as e:
            logger.error(f"发送Matrix消息失败: {e}")
            return False
    
    async def disconnect(self):
        """断开连接"""
        if self.client:
            try:
                await self.client.logout()
                await self.client.close()
            except:
                pass
            self.is_running = False