加入新功能:番号查询。

This commit is contained in:
liuwei
2025-09-10 16:38:04 +08:00
parent 11b4144681
commit 55cbf222cb

View File

@@ -1,4 +1,5 @@
from typing import Dict, Any, List, Optional, Tuple
import re
from loguru import logger
from base.plugin_common.message_plugin_interface import MessagePluginInterface
@@ -100,35 +101,71 @@ class FanhaoSearchPlugin(MessagePluginInterface):
command = content.split(" ")[0]
return command in self._commands
def _redact_mongo_uri(self, uri: str) -> str:
try:
# 隐藏用户名密码,仅保留协议和主机段
return re.sub(r"(mongodb\+srv://)(.*?@)", r"\\1***@", uri)
except Exception:
return "***"
def _ensure_mongo(self):
if self.mongo_client:
return
from pymongo import MongoClient
# 采用默认匿名直连如需凭证可在URI中配置
self.mongo_client = MongoClient(self.mongo_uri)
self.mongo_db = self.mongo_client.get_database(self.mongo_db_name)
self.LOG.info(
f"[{self.name}] 准备连接MongoDB uri={self._redact_mongo_uri(self.mongo_uri)} db={self.mongo_db_name}"
)
try:
self.mongo_client = MongoClient(self.mongo_uri, serverSelectionTimeoutMS=5000)
# 探活
self.mongo_client.admin.command("ping")
self.mongo_db = self.mongo_client.get_database(self.mongo_db_name)
try:
colls = self.mongo_db.list_collection_names()
except Exception as e:
colls = []
self.LOG.warning(f"[{self.name}] 获取集合列表失败: {e}")
self.LOG.info(f"[{self.name}] MongoDB连接成功集合={colls}")
except Exception as e:
self.LOG.error(f"[{self.name}] MongoDB连接失败: {e}")
raise
def _normalize_code(self, text: str) -> str:
return (text or "").strip().upper()
def _build_queries(self, code_upper: str) -> List[Dict[str, Any]]:
# 精确匹配查询
or_exact = [{field: code_upper} for field in self.search_fields]
exact_query = {"$or": or_exact}
# 回退:大小写不敏感的等值匹配
or_regex = [
{field: {"$regex": f"^{re.escape(code_upper)}$", "$options": "i"}}
for field in self.search_fields
]
regex_query = {"$or": or_regex}
return [exact_query, regex_query]
def _query_collections(self, code_upper: str) -> Optional[Dict[str, Any]]:
self._ensure_mongo()
# 构造或条件:任一字段等于 code_upper
or_conditions = [{field: code_upper} for field in self.search_fields]
query = {"$or": or_conditions}
for coll_name in self.collections:
try:
coll = self.mongo_db.get_collection(coll_name)
doc = coll.find_one(query)
if doc:
doc["_collection"] = coll_name
return doc
except Exception as e:
self.LOG.error(f"查询集合 {coll_name} 出错: {e}")
continue
queries = self._build_queries(code_upper)
self.LOG.debug(f"[{self.name}] 标准化番号={code_upper},查询字段={self.search_fields}")
for idx, query in enumerate(queries):
self.LOG.debug(f"[{self.name}] 执行查询({idx+1}/{len(queries)}): {query}")
for coll_name in self.collections:
try:
coll = self.mongo_db.get_collection(coll_name)
self.LOG.debug(f"[{self.name}] 在集合 {coll_name} 查找…")
doc = coll.find_one(query)
if doc:
doc["_collection"] = coll_name
self.LOG.info(f"[{self.name}] 命中集合 {coll_name}")
return doc
else:
self.LOG.debug(f"[{self.name}] 集合 {coll_name} 未命中")
except Exception as e:
self.LOG.error(f"[{self.name}] 查询集合 {coll_name} 出错: {e}")
continue
return None
def _format_result(self, doc: Dict[str, Any]) -> str:
@@ -190,7 +227,11 @@ class FanhaoSearchPlugin(MessagePluginInterface):
if roomid and gbm.get_group_permission(roomid, self.feature) == PermissionStatus.DISABLED:
return False, "没有权限"
user_code = self._normalize_code(content[len(command):].strip())
raw_code = content[len(command):].strip()
user_code = self._normalize_code(raw_code)
self.LOG.info(
f"[{self.name}] 收到查询 command={command} raw='{raw_code}' normalized='{user_code}' db={self.mongo_db_name} collections={self.collections}"
)
if not user_code:
await bot.send_text_message((roomid if roomid else sender), f"❌命令格式错误!\n{self.command_format}", sender)
return False, "命令格式错误"
@@ -199,6 +240,7 @@ class FanhaoSearchPlugin(MessagePluginInterface):
doc = self._query_collections(user_code)
target = roomid if roomid else sender
if not doc:
self.LOG.warning(f"[{self.name}] 未找到番号:{user_code}")
await bot.send_text_message(target, f"未找到番号:{user_code}", sender)
return False, "未找到"
@@ -206,7 +248,7 @@ class FanhaoSearchPlugin(MessagePluginInterface):
await bot.send_text_message(target, text, sender)
return True, "查询成功"
except Exception as e:
self.LOG.error(f"处理番号查询出错: {e}")
self.LOG.exception(f"处理番号查询出错: {e}")
return False, f"处理出错: {e}"