208 lines
7.2 KiB
Python
208 lines
7.2 KiB
Python
"""
|
||
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
|
||
|
||
|
||
# 全局实例
|
||
matrix_bot = MatrixBot() |