diff --git a/app/api/v1/settings.py b/app/api/v1/settings.py index 94e852a..1cd51d9 100644 --- a/app/api/v1/settings.py +++ b/app/api/v1/settings.py @@ -12,6 +12,15 @@ from app.config import settings router = APIRouter() +def _mask_secret(value: str | None, visible: int = 4) -> str | None: + """对敏感字符串脱敏，保留前 visible 个字符""" + if not value: + return None + if len(value) <= visible: + return "****" + return value[:visible] + "****" + + class SettingsResponse(BaseModel): """设置响应（隐藏敏感信息的中间位）""" embedding_provider: str @@ -37,6 +46,10 @@ class SettingsResponse(BaseModel): has_dashscope_key: bool has_aliyun_ocr: bool has_tencent_ocr: bool + # 企业微信机器人 + wework_bot_enabled: bool + wework_bot_id: str | None + wework_bot_secret_masked: str | None @router.get("", response_model=SettingsResponse, summary="获取当前设置") @@ -65,6 +78,9 @@ async def get_settings(): has_dashscope_key=bool(settings.DASHSCOPE_API_KEY), has_aliyun_ocr=bool(settings.ALIYUN_OCR_ACCESS_KEY), has_tencent_ocr=bool(settings.TENCENT_OCR_SECRET_ID), + wework_bot_enabled=settings.WEWORK_BOT_ENABLED, + wework_bot_id=settings.WEWORK_BOT_ID, + wework_bot_secret_masked=_mask_secret(settings.WEWORK_BOT_SECRET), ) @@ -94,6 +110,10 @@ class SettingsUpdate(BaseModel): chunk_overlap: int | None = None search_limit: int | None = None judge_batch_size: int | None = None + # 企业微信机器人 + wework_bot_enabled: bool | None = None + wework_bot_id: str | None = None + wework_bot_secret: str | None = None @router.put("", summary="更新设置") @@ -139,6 +159,9 @@ async def update_settings(data: SettingsUpdate): "chunk_overlap": "CHUNK_OVERLAP", "search_limit": "SEARCH_LIMIT", "judge_batch_size": "JUDGE_BATCH_SIZE", + "wework_bot_enabled": "WEWORK_BOT_ENABLED", + "wework_bot_id": "WEWORK_BOT_ID", + "wework_bot_secret": "WEWORK_BOT_SECRET", } for api_field, config_field in field_map.items(): @@ -221,3 +244,18 @@ async def test_ocr(): return {"success": True, "message": f"OCR 提供商 {provider.name} 已就绪"} except Exception as e: return {"success": False, "message": f"OCR 测试失败: {str(e)}"} + + +@router.post("/bot/restart", summary="重启企业微信机器人") +async def restart_bot(): + """重启企业微信机器人服务""" + from app.main import bot_manager + result = await bot_manager.restart() + return result + + +@router.get("/bot/status", summary="获取企业微信机器人状态") +async def get_bot_status(): + """获取企业微信机器人运行状态""" + from app.main import bot_manager + return bot_manager.get_status() diff --git a/app/main.py b/app/main.py index fd750f2..08dc202 100644 --- a/app/main.py +++ b/app/main.py @@ -5,8 +5,10 @@ FastAPI 应用入口 from __future__ import annotations +import asyncio import json import logging +import threading from contextlib import asynccontextmanager from pathlib import Path from typing import AsyncGenerator @@ -18,6 +20,88 @@ from fastapi.responses import FileResponse, JSONResponse, RedirectResponse from app.config import settings from app.database import close_db, init_db +_logger = logging.getLogger(__name__) + + +class BotManager: + """管理企业微信机器人的生命周期（在后台线程中运行）""" + + def __init__(self): + self._thread: threading.Thread | None = None + self._running = False + + def start(self): + """在后台线程中启动机器人""" + if self._running: + _logger.info("企业微信机器人已在运行") + return + if not settings.WEWORK_BOT_ENABLED or not settings.WEWORK_BOT_ID or not settings.WEWORK_BOT_SECRET: + _logger.info("企业微信机器人未配置，跳过启动") + return + + self._thread = threading.Thread( + target=self._run_bot, daemon=True, name="wework-bot" + ) + self._thread.start() + _logger.info("企业微信机器人已在后台线程启动") + + def _run_bot(self): + """在线程中运行 bot（需要新建事件循环）""" + self._running = True + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + loop.run_until_complete(self._run_bot_async()) + except Exception as e: + _logger.error("企业微信机器人线程异常退出: %s", e, exc_info=True) + finally: + loop.close() + self._running = False + + async def _run_bot_async(self): + """异步运行 bot""" + from app.wework_bot import WeComBotService + service = WeComBotService() + await service.start_async() + + async def restart(self) -> dict: + """重启机器人（先停后启）""" + was_running = self._running + # 停止旧的（设置 enabled=False 让旧线程自然退出） + old_id = settings.WEWORK_BOT_ID + old_secret = settings.WEWORK_BOT_SECRET + old_enabled = settings.WEWORK_BOT_ENABLED + + if not old_enabled or not old_id or not old_secret: + return {"success": False, "message": "请先配置机器人 ID 和 Secret 并启用"} + + # 等待旧线程结束 + if self._thread and self._thread.is_alive(): + # SDK 的 run() 会阻塞，无法从外部优雅停止 + # 通过设置 enabled=False 让下次检查时退出 + _logger.info("等待旧机器人线程结束...") + # 无法强制停止，提示用户 + if was_running: + return {"success": False, "message": "旧机器人正在运行，请稍后重试（SDK 不支持强制停止）"} + + # 启动新的 + self.start() + return {"success": True, "message": "机器人已启动"} + + def get_status(self) -> dict: + """获取机器人状态""" + configured = bool(settings.WEWORK_BOT_ID and settings.WEWORK_BOT_SECRET) + return { + "enabled": settings.WEWORK_BOT_ENABLED, + "configured": configured, + "running": self._running, + "bot_id": settings.WEWORK_BOT_ID, + } + + +# 全局 bot 管理器 +bot_manager = BotManager() + @asynccontextmanager async def lifespan(_app: FastAPI) -> AsyncGenerator[None, None]: @@ -32,6 +116,9 @@ async def lifespan(_app: FastAPI) -> AsyncGenerator[None, None]: settings.ensure_dirs() await init_db() + # 启动企业微信机器人（后台线程） + bot_manager.start() + # 预热嵌入服务（懒加载，首次调用时初始化） try: from app.services.embedding_service import EmbeddingService diff --git a/app/wework_bot.py b/app/wework_bot.py index fcd63c8..cba3324 100644 --- a/app/wework_bot.py +++ b/app/wework_bot.py @@ -18,6 +18,7 @@ import os import re import tempfile import time +import asyncio from datetime import datetime import httpx @@ -47,7 +48,7 @@ class WeComBotService: self.ws_client: WSClient | None = None def start(self): - """启动机器人服务""" + """启动机器人服务（阻塞模式，用于独立进程运行）""" if not self.enabled: logger.info("企业微信机器人未启用 (WEWORK_BOT_ENABLED=false)") return @@ -63,7 +64,35 @@ class WeComBotService: ) ) - # 注册事件 + self._register_events() + logger.info("启动企业微信智能机器人...") + self.ws_client.run() + + async def start_async(self): + """启动机器人服务（异步模式，用于嵌入到其他事件循环）""" + if not self.enabled: + logger.info("企业微信机器人未启用 (WEWORK_BOT_ENABLED=false)") + return + + if not self.bot_id or not self.bot_secret: + logger.error("WEWORK_BOT_ID 或 WEWORK_BOT_SECRET 未配置") + return + + self.ws_client = WSClient( + WSClientOptions( + bot_id=self.bot_id, + secret=self.bot_secret, + ) + ) + + self._register_events() + logger.info("启动企业微信智能机器人（异步模式）...") + await self.ws_client.connect() + # 永远等待，保持连接 + await asyncio.Future() + + def _register_events(self): + """注册所有事件处理器""" self.ws_client.on("authenticated", self._on_authenticated) self.ws_client.on("message.text", self._on_text_message) self.ws_client.on("message.image", self._on_image_message) @@ -71,9 +100,6 @@ class WeComBotService: self.ws_client.on("disconnected", self._on_disconnected) self.ws_client.on("error", self._on_error) - logger.info("启动企业微信智能机器人...") - self.ws_client.run() - # ──────────────────────────── 事件处理 ──────────────────────────── def _on_authenticated(self): diff --git a/static/index.html b/static/index.html index 8cfc078..d45e570 100644 --- a/static/index.html +++ b/static/index.html @@ -1041,6 +1041,38 @@ + +
+
+

企业微信机器人

+
+ 检测中 + +
+
+
+
+ + +
+
+ + +
+
+ + +
+
+ + +
+
+
+

图片搜索设置

@@ -2612,12 +2644,22 @@ if (data.judge_batch_size != null) { document.getElementById('setting-judge-batch-size').value = data.judge_batch_size; } + // 企业微信机器人 + document.getElementById('setting-wework-bot-enabled').checked = data.wework_bot_enabled; + document.getElementById('setting-wework-bot-id').value = data.wework_bot_id || ''; + // secret 只显示脱敏值作为 placeholder + const secretInput = document.getElementById('setting-wework-bot-secret'); + if (data.wework_bot_secret_masked) { + secretInput.placeholder = '当前: ' + data.wework_bot_secret_masked; + } } catch (_) { // 静默处理 } // 初始化字段显示状态 toggleEmbedFields(); toggleOcrFields(); + // 加载机器人状态 + loadBotStatus(); } /** 保存设置 */ @@ -2631,6 +2673,12 @@ if (!isNaN(judgeBatchSize) && judgeBatchSize >= 1 && judgeBatchSize <= 50) { body.judge_batch_size = judgeBatchSize; } + // 企业微信机器人 + body.wework_bot_enabled = document.getElementById('setting-wework-bot-enabled').checked; + const botId = document.getElementById('setting-wework-bot-id').value.trim(); + if (botId) body.wework_bot_id = botId; + const botSecret = document.getElementById('setting-wework-bot-secret').value.trim(); + if (botSecret) body.wework_bot_secret = botSecret; if (Object.keys(body).length === 0) { showToast('没有需要保存的更改', 'info'); return; @@ -2643,6 +2691,53 @@ } } + /** 加载机器人状态 */ + async function loadBotStatus() { + try { + const data = await api('GET', '/settings/bot/status'); + const badge = document.getElementById('bot-status-badge'); + if (data.running) { + badge.className = 'badge badge-completed'; + badge.textContent = '运行中'; + } else if (data.configured) { + badge.className = 'badge badge-failed'; + badge.textContent = '已停止'; + } else { + badge.className = 'badge badge-pending'; + badge.textContent = '未配置'; + } + } catch (_) { + document.getElementById('bot-status-badge').className = 'badge badge-pending'; + document.getElementById('bot-status-badge').textContent = '未知'; + } + } + + /** 重启企业微信机器人 */ + async function restartBot() { + const btn = document.getElementById('bot-restart-btn'); + btn.disabled = true; + btn.textContent = '重启中...'; + try { + const result = await api('POST', '/settings/bot/restart'); + if (result.success) { + showToast(result.message, 'success'); + } else { + showToast(result.message, 'error'); + } + setTimeout(() => loadBotStatus(), 2000); + } catch (err) { + showToast('重启失败: ' + err.message, 'error'); + } finally { + btn.disabled = false; + btn.textContent = '重启机器人'; + } + } + + /** 切换机器人启用状态 */ + function toggleBotEnabled() { + // 仅切换 UI 状态，实际保存需要点"保存设置" + } + /** 测试嵌入模型连接 */ async function testEmbedding() { showToast('正在测试嵌入模型连接...', 'info');