dbassistant.py 65 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848
  1. from __future__ import annotations
  2. import asyncio
  3. import hashlib
  4. import json
  5. import re
  6. from datetime import UTC, datetime, timedelta
  7. from typing import Any
  8. from urllib.parse import urlsplit
  9. from uuid import uuid4
  10. from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
  11. from pymongo import ASCENDING, DESCENDING
  12. from pymongo.errors import DuplicateKeyError
  13. from wbb import BOT_PROFILE_ID, db
  14. connectionsdb = db.business_assistant_connections
  15. settingsdb = db.business_assistant_settings
  16. knowledgedb = db.business_assistant_knowledge
  17. knowledge_sourcesdb = db.business_assistant_knowledge_sources
  18. source_eventsdb = db.business_assistant_source_events
  19. knowledge_candidatesdb = db.business_assistant_knowledge_candidates
  20. conversationsdb = db.business_assistant_conversations
  21. messagesdb = db.business_assistant_messages
  22. usagedb = db.business_assistant_usage
  23. updatesdb = db.business_assistant_updates
  24. deadlettersdb = db.business_assistant_dead_letters
  25. runtimedb = db.business_assistant_runtime
  26. DEFAULT_ACCOUNT_SETTINGS: dict[str, Any] = {
  27. "assistant_enabled": False,
  28. "system_prompt": (
  29. "你是该 Telegram 账号的智能接待秘书。回答要简洁、礼貌,只能依据提供的知识条目陈述业务事实。"
  30. ),
  31. "language": "zh-CN",
  32. "tone": "professional",
  33. "account_daily_limit": 200,
  34. "customer_daily_limit": 20,
  35. "human_pause_hours": 24,
  36. "notification_destination": "owner",
  37. "ops_group_id": 0,
  38. "timezone": "Asia/Shanghai",
  39. "digest_enabled": False,
  40. "digest_time": "09:00",
  41. "handoff_message": "这个问题需要人工确认,我已经通知负责人,请稍候。",
  42. "unsupported_message": "已收到你的消息,这类内容需要人工处理,我已经通知负责人。",
  43. "feishu_webhook_enabled": False,
  44. "feishu_webhook_url": "",
  45. "feishu_webhook_signing_secret": "",
  46. "feishu_message_preview_enabled": False,
  47. "feishu_outgoing_message_enabled": False,
  48. }
  49. SECRET_ACCOUNT_SETTINGS = {
  50. "feishu_webhook_url",
  51. "feishu_webhook_signing_secret",
  52. }
  53. VALID_NOTIFICATION_DESTINATIONS = {"owner", "ops", "both"}
  54. VALID_TONES = {"professional", "friendly", "concise"}
  55. VALID_SOURCE_TYPES = {"channel", "group", "business"}
  56. VALID_PUBLICATION_MODES = {"review", "auto"}
  57. VALID_AUTHOR_POLICIES = {"admins_or_allowlist", "admins", "allowlist", "all"}
  58. VALID_CANDIDATE_STATUSES = {"pending", "published", "rejected", "stale"}
  59. HANDOFF_TERMS = (
  60. "人工",
  61. "真人",
  62. "客服",
  63. "负责人",
  64. "转人工",
  65. "human",
  66. "agent",
  67. )
  68. SENSITIVE_TERMS = (
  69. "承诺",
  70. "保证",
  71. "投诉",
  72. "退款",
  73. "退钱",
  74. "赔偿",
  75. "律师",
  76. "起诉",
  77. "支付失败",
  78. "账号被盗",
  79. "密码",
  80. "验证码",
  81. )
  82. _index_lock = asyncio.Lock()
  83. _indexes_ready = False
  84. class AssistantDataError(ValueError):
  85. def __init__(self, code: str, message: str):
  86. super().__init__(message)
  87. self.code = code
  88. def utc_now() -> datetime:
  89. return datetime.now(UTC)
  90. def as_utc(value: datetime) -> datetime:
  91. return value.replace(tzinfo=UTC) if value.tzinfo is None else value.astimezone(UTC)
  92. def _scope(filters: dict[str, Any] | None = None, /, **values: Any) -> dict[str, Any]:
  93. return {"bot_id": BOT_PROFILE_ID, **(filters or {}), **values}
  94. def clean_text(
  95. value: Any,
  96. *,
  97. max_length: int,
  98. required: bool = False,
  99. preserve_lines: bool = False,
  100. ) -> str:
  101. text = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]", "", str(value or ""))
  102. text = text.strip()
  103. if not preserve_lines:
  104. text = " ".join(text.split())
  105. if required and not text:
  106. raise AssistantDataError("required_field", "必填内容不能为空。")
  107. if len(text) > max_length:
  108. raise AssistantDataError("text_too_long", f"内容不能超过 {max_length} 个字符。")
  109. return text
  110. def _bounded_int(
  111. value: Any,
  112. name: str,
  113. *,
  114. minimum: int,
  115. maximum: int,
  116. ) -> int:
  117. try:
  118. parsed = int(value)
  119. except (TypeError, ValueError) as exc:
  120. raise AssistantDataError("invalid_setting", f"{name} 必须是整数。") from exc
  121. if not minimum <= parsed <= maximum:
  122. raise AssistantDataError(
  123. "invalid_setting",
  124. f"{name} 必须在 {minimum} 到 {maximum} 之间。",
  125. )
  126. return parsed
  127. def _normalize_feishu_webhook_url(value: Any) -> str:
  128. url = clean_text(value, max_length=1000)
  129. if not url:
  130. return ""
  131. try:
  132. parsed = urlsplit(url)
  133. port = parsed.port
  134. except ValueError as exc:
  135. raise AssistantDataError("invalid_setting", "飞书 Webhook 地址无效。") from exc
  136. allowed_hosts = {"open.feishu.cn", "open.larksuite.com"}
  137. if (
  138. parsed.scheme != "https"
  139. or (parsed.hostname or "").casefold() not in allowed_hosts
  140. or port not in {None, 443}
  141. or parsed.username
  142. or parsed.password
  143. or parsed.query
  144. or parsed.fragment
  145. or not re.fullmatch(r"/open-apis/bot/v2/hook/[A-Za-z0-9-]{8,128}", parsed.path)
  146. ):
  147. raise AssistantDataError(
  148. "invalid_setting",
  149. "仅支持飞书或 Lark 官方 HTTPS 群机器人 Webhook 地址。",
  150. )
  151. return url
  152. def _normalize_string_list(value: Any, *, max_items: int, max_length: int) -> list[str]:
  153. values = value if isinstance(value, (list, tuple, set)) else str(value or "").split(",")
  154. normalized: list[str] = []
  155. seen: set[str] = set()
  156. for item in values:
  157. text = clean_text(item, max_length=max_length)
  158. key = text.casefold()
  159. if text and key not in seen:
  160. normalized.append(text)
  161. seen.add(key)
  162. if len(normalized) > max_items:
  163. raise AssistantDataError("too_many_items", f"最多允许 {max_items} 项。")
  164. return normalized
  165. def _normalize_int_list(value: Any, *, max_items: int = 100) -> list[int]:
  166. values = value if isinstance(value, (list, tuple, set)) else str(value or "").split(",")
  167. normalized: list[int] = []
  168. for item in values:
  169. if item in {None, ""}:
  170. continue
  171. try:
  172. parsed = int(item)
  173. except (TypeError, ValueError) as exc:
  174. raise AssistantDataError("invalid_source", "白名单用户 ID 必须是整数。") from exc
  175. if parsed and parsed not in normalized:
  176. normalized.append(parsed)
  177. if len(normalized) > max_items:
  178. raise AssistantDataError("too_many_items", f"最多允许 {max_items} 项。")
  179. return normalized
  180. async def ensure_assistant_indexes() -> None:
  181. global _indexes_ready
  182. if _indexes_ready:
  183. return
  184. async with _index_lock:
  185. if _indexes_ready:
  186. return
  187. await connectionsdb.create_index(
  188. [("bot_id", ASCENDING), ("connection_id", ASCENDING)], unique=True
  189. )
  190. await connectionsdb.create_index(
  191. [("bot_id", ASCENDING), ("updated_at", DESCENDING)]
  192. )
  193. await settingsdb.create_index(
  194. [("bot_id", ASCENDING), ("connection_id", ASCENDING)], unique=True
  195. )
  196. await knowledgedb.create_index(
  197. [("bot_id", ASCENDING), ("entry_id", ASCENDING)], unique=True
  198. )
  199. await knowledge_sourcesdb.create_index(
  200. [("bot_id", ASCENDING), ("source_id", ASCENDING)], unique=True
  201. )
  202. await knowledge_sourcesdb.create_index(
  203. [("bot_id", ASCENDING), ("source_key", ASCENDING)], unique=True
  204. )
  205. await knowledge_sourcesdb.create_index(
  206. [
  207. ("bot_id", ASCENDING),
  208. ("connection_id", ASCENDING),
  209. ("enabled", ASCENDING),
  210. ("updated_at", DESCENDING),
  211. ]
  212. )
  213. await source_eventsdb.create_index(
  214. [("bot_id", ASCENDING), ("event_id", ASCENDING)], unique=True
  215. )
  216. await source_eventsdb.create_index(
  217. [
  218. ("bot_id", ASCENDING),
  219. ("source_id", ASCENDING),
  220. ("event_key", ASCENDING),
  221. ],
  222. unique=True,
  223. )
  224. await source_eventsdb.create_index("expires_at", expireAfterSeconds=0)
  225. await knowledge_candidatesdb.create_index(
  226. [("bot_id", ASCENDING), ("candidate_id", ASCENDING)], unique=True
  227. )
  228. await knowledge_candidatesdb.create_index(
  229. [
  230. ("bot_id", ASCENDING),
  231. ("event_id", ASCENDING),
  232. ("item_index", ASCENDING),
  233. ],
  234. unique=True,
  235. )
  236. await knowledge_candidatesdb.create_index(
  237. [
  238. ("bot_id", ASCENDING),
  239. ("connection_id", ASCENDING),
  240. ("status", ASCENDING),
  241. ("updated_at", DESCENDING),
  242. ]
  243. )
  244. await knowledgedb.create_index(
  245. [
  246. ("bot_id", ASCENDING),
  247. ("connection_id", ASCENDING),
  248. ("enabled", ASCENDING),
  249. ("priority", DESCENDING),
  250. ]
  251. )
  252. await conversationsdb.create_index(
  253. [("bot_id", ASCENDING), ("conversation_id", ASCENDING)], unique=True
  254. )
  255. await conversationsdb.create_index(
  256. [
  257. ("bot_id", ASCENDING),
  258. ("connection_id", ASCENDING),
  259. ("chat_id", ASCENDING),
  260. ],
  261. unique=True,
  262. )
  263. await conversationsdb.create_index(
  264. [
  265. ("bot_id", ASCENDING),
  266. ("connection_id", ASCENDING),
  267. ("status", ASCENDING),
  268. ("updated_at", DESCENDING),
  269. ]
  270. )
  271. await messagesdb.create_index(
  272. [("bot_id", ASCENDING), ("message_key", ASCENDING)], unique=True
  273. )
  274. await messagesdb.create_index(
  275. [
  276. ("bot_id", ASCENDING),
  277. ("conversation_id", ASCENDING),
  278. ("created_at", ASCENDING),
  279. ]
  280. )
  281. await messagesdb.create_index("expires_at", expireAfterSeconds=0)
  282. await usagedb.create_index(
  283. [
  284. ("bot_id", ASCENDING),
  285. ("connection_id", ASCENDING),
  286. ("chat_id", ASCENDING),
  287. ("day", ASCENDING),
  288. ],
  289. unique=True,
  290. )
  291. await updatesdb.create_index(
  292. [("bot_id", ASCENDING), ("update_id", ASCENDING)], unique=True
  293. )
  294. await updatesdb.create_index("expires_at", expireAfterSeconds=0)
  295. await deadlettersdb.create_index(
  296. [("bot_id", ASCENDING), ("update_id", ASCENDING)], unique=True
  297. )
  298. await deadlettersdb.create_index(
  299. [("bot_id", ASCENDING), ("created_at", DESCENDING)]
  300. )
  301. await deadlettersdb.create_index("expires_at", expireAfterSeconds=0)
  302. await runtimedb.create_index(
  303. [("bot_id", ASCENDING), ("runtime_id", ASCENDING)], unique=True
  304. )
  305. _indexes_ready = True
  306. def _public_user(user: Any) -> dict[str, Any]:
  307. value = user if isinstance(user, dict) else {}
  308. return {
  309. "id": int(value.get("id") or 0),
  310. "username": clean_text(value.get("username"), max_length=64),
  311. "first_name": clean_text(value.get("first_name"), max_length=128),
  312. "last_name": clean_text(value.get("last_name"), max_length=128),
  313. }
  314. async def upsert_business_connection(payload: dict[str, Any]) -> dict[str, Any]:
  315. await ensure_assistant_indexes()
  316. connection_id = clean_text(payload.get("id"), max_length=128, required=True)
  317. now = utc_now()
  318. rights = payload.get("rights") if isinstance(payload.get("rights"), dict) else {}
  319. document = {
  320. "bot_id": BOT_PROFILE_ID,
  321. "connection_id": connection_id,
  322. "user": _public_user(payload.get("user")),
  323. "user_chat_id": int(payload.get("user_chat_id") or 0),
  324. "rights": {str(key): bool(value) for key, value in rights.items() if value is True},
  325. "is_enabled": bool(payload.get("is_enabled")),
  326. "connected_at": datetime.fromtimestamp(int(payload.get("date") or 0), UTC)
  327. if payload.get("date")
  328. else now,
  329. "last_event_at": now,
  330. "updated_at": now,
  331. }
  332. await connectionsdb.update_one(
  333. _scope(connection_id=connection_id),
  334. {"$set": document, "$setOnInsert": {"created_at": now}},
  335. upsert=True,
  336. )
  337. await settingsdb.update_one(
  338. _scope(connection_id=connection_id),
  339. {
  340. "$setOnInsert": {
  341. "bot_id": BOT_PROFILE_ID,
  342. "connection_id": connection_id,
  343. **DEFAULT_ACCOUNT_SETTINGS,
  344. "created_at": now,
  345. "updated_at": now,
  346. }
  347. },
  348. upsert=True,
  349. )
  350. return await get_business_connection(connection_id) or document
  351. async def get_business_connection(connection_id: str) -> dict[str, Any] | None:
  352. await ensure_assistant_indexes()
  353. return await connectionsdb.find_one(_scope(connection_id=str(connection_id)))
  354. async def touch_business_connection(
  355. connection_id: str, *, error: str = "", is_enabled: bool | None = None
  356. ) -> dict[str, Any] | None:
  357. now = utc_now()
  358. values: dict[str, Any] = {
  359. "last_event_at": now,
  360. "updated_at": now,
  361. "last_error": clean_text(error, max_length=1000),
  362. }
  363. if is_enabled is not None:
  364. values["is_enabled"] = bool(is_enabled)
  365. await connectionsdb.update_one(
  366. _scope(connection_id=str(connection_id)),
  367. {"$set": values},
  368. )
  369. return await get_business_connection(connection_id)
  370. async def list_business_connections(
  371. *, page: int = 1, page_size: int = 20
  372. ) -> tuple[list[dict[str, Any]], int]:
  373. await ensure_assistant_indexes()
  374. page = max(1, int(page))
  375. page_size = max(1, min(int(page_size), 100))
  376. filters = _scope()
  377. total = await connectionsdb.count_documents(filters)
  378. cursor = (
  379. connectionsdb.find(filters)
  380. .sort("updated_at", DESCENDING)
  381. .skip((page - 1) * page_size)
  382. .limit(page_size)
  383. )
  384. items = []
  385. async for item in cursor:
  386. item["settings"] = public_account_settings(
  387. await get_account_settings(item["connection_id"])
  388. )
  389. items.append(item)
  390. return items, total
  391. async def get_account_settings(connection_id: str) -> dict[str, Any]:
  392. await ensure_assistant_indexes()
  393. stored = await settingsdb.find_one(_scope(connection_id=str(connection_id))) or {}
  394. return {
  395. **DEFAULT_ACCOUNT_SETTINGS,
  396. **{key: value for key, value in stored.items() if key != "_id"},
  397. "connection_id": str(connection_id),
  398. }
  399. def public_account_settings(settings: dict[str, Any]) -> dict[str, Any]:
  400. public = {
  401. key: value
  402. for key, value in settings.items()
  403. if key not in SECRET_ACCOUNT_SETTINGS
  404. }
  405. public.update(
  406. {
  407. "feishu_webhook_url": "",
  408. "feishu_webhook_signing_secret": "",
  409. "feishu_webhook_configured": bool(settings.get("feishu_webhook_url")),
  410. "feishu_signing_secret_configured": bool(
  411. settings.get("feishu_webhook_signing_secret")
  412. ),
  413. }
  414. )
  415. return public
  416. def normalize_account_settings(
  417. values: dict[str, Any], *, previous: dict[str, Any] | None = None
  418. ) -> dict[str, Any]:
  419. current = {**DEFAULT_ACCOUNT_SETTINGS, **(previous or {})}
  420. if "assistant_enabled" in values:
  421. current["assistant_enabled"] = bool(values.get("assistant_enabled"))
  422. if "system_prompt" in values:
  423. current["system_prompt"] = clean_text(
  424. values.get("system_prompt"), max_length=4000, required=True, preserve_lines=True
  425. )
  426. if "language" in values:
  427. current["language"] = clean_text(values.get("language"), max_length=32, required=True)
  428. if "tone" in values:
  429. tone = str(values.get("tone") or "")
  430. if tone not in VALID_TONES:
  431. raise AssistantDataError("invalid_setting", "接待语气无效。")
  432. current["tone"] = tone
  433. if "account_daily_limit" in values:
  434. current["account_daily_limit"] = _bounded_int(
  435. values.get("account_daily_limit"), "账号每日额度", minimum=1, maximum=100000
  436. )
  437. if "customer_daily_limit" in values:
  438. current["customer_daily_limit"] = _bounded_int(
  439. values.get("customer_daily_limit"), "客户每日额度", minimum=1, maximum=1000
  440. )
  441. if "human_pause_hours" in values:
  442. current["human_pause_hours"] = _bounded_int(
  443. values.get("human_pause_hours"), "人工暂停小时数", minimum=1, maximum=168
  444. )
  445. if "notification_destination" in values:
  446. destination = str(values.get("notification_destination") or "")
  447. if destination not in VALID_NOTIFICATION_DESTINATIONS:
  448. raise AssistantDataError("invalid_setting", "通知目的地无效。")
  449. current["notification_destination"] = destination
  450. if "ops_group_id" in values:
  451. try:
  452. current["ops_group_id"] = int(values.get("ops_group_id") or 0)
  453. except (TypeError, ValueError) as exc:
  454. raise AssistantDataError("invalid_setting", "运营群 ID 必须是整数。") from exc
  455. if "timezone" in values:
  456. timezone_name = clean_text(values.get("timezone"), max_length=64, required=True)
  457. try:
  458. ZoneInfo(timezone_name)
  459. except ZoneInfoNotFoundError as exc:
  460. raise AssistantDataError("invalid_setting", "时区名称无效。") from exc
  461. current["timezone"] = timezone_name
  462. if "digest_enabled" in values:
  463. current["digest_enabled"] = bool(values.get("digest_enabled"))
  464. if "digest_time" in values:
  465. digest_time = str(values.get("digest_time") or "")
  466. if not re.fullmatch(r"(?:[01]\d|2[0-3]):[0-5]\d", digest_time):
  467. raise AssistantDataError("invalid_setting", "简报时间必须是 HH:MM。")
  468. current["digest_time"] = digest_time
  469. for key, max_length in (("handoff_message", 1000), ("unsupported_message", 1000)):
  470. if key in values:
  471. current[key] = clean_text(values.get(key), max_length=max_length, required=True)
  472. if "feishu_webhook_url" in values:
  473. current["feishu_webhook_url"] = _normalize_feishu_webhook_url(
  474. values.get("feishu_webhook_url")
  475. )
  476. if "feishu_webhook_signing_secret" in values:
  477. current["feishu_webhook_signing_secret"] = clean_text(
  478. values.get("feishu_webhook_signing_secret"), max_length=512
  479. )
  480. if "feishu_webhook_enabled" in values:
  481. current["feishu_webhook_enabled"] = bool(values.get("feishu_webhook_enabled"))
  482. if "feishu_message_preview_enabled" in values:
  483. current["feishu_message_preview_enabled"] = bool(
  484. values.get("feishu_message_preview_enabled")
  485. )
  486. if "feishu_outgoing_message_enabled" in values:
  487. current["feishu_outgoing_message_enabled"] = bool(
  488. values.get("feishu_outgoing_message_enabled")
  489. )
  490. if current["feishu_webhook_enabled"] and not current["feishu_webhook_url"]:
  491. raise AssistantDataError(
  492. "invalid_setting", "开启飞书消息提醒前,请先填写群机器人 Webhook 地址。"
  493. )
  494. return {key: current[key] for key in DEFAULT_ACCOUNT_SETTINGS}
  495. async def update_account_settings(
  496. connection_id: str, values: dict[str, Any]
  497. ) -> dict[str, Any]:
  498. connection = await get_business_connection(connection_id)
  499. if not connection:
  500. raise AssistantDataError("connection_not_found", "未找到 Business 连接。")
  501. previous = await get_account_settings(connection_id)
  502. normalized = normalize_account_settings(values, previous=previous)
  503. now = utc_now()
  504. await settingsdb.update_one(
  505. _scope(connection_id=str(connection_id)),
  506. {
  507. "$set": {**normalized, "updated_at": now},
  508. "$setOnInsert": {
  509. "bot_id": BOT_PROFILE_ID,
  510. "connection_id": str(connection_id),
  511. "created_at": now,
  512. },
  513. },
  514. upsert=True,
  515. )
  516. return await get_account_settings(connection_id)
  517. async def create_knowledge_entry(
  518. connection_id: str, values: dict[str, Any]
  519. ) -> dict[str, Any]:
  520. if not await get_business_connection(connection_id):
  521. raise AssistantDataError("connection_not_found", "未找到 Business 连接。")
  522. now = utc_now()
  523. entry = {
  524. "bot_id": BOT_PROFILE_ID,
  525. "entry_id": uuid4().hex,
  526. "connection_id": str(connection_id),
  527. "question": clean_text(values.get("question"), max_length=300, required=True),
  528. "aliases": _normalize_string_list(values.get("aliases"), max_items=20, max_length=200),
  529. "keywords": _normalize_string_list(values.get("keywords"), max_items=30, max_length=50),
  530. "answer": clean_text(
  531. values.get("answer"), max_length=4000, required=True, preserve_lines=True
  532. ),
  533. "tags": _normalize_string_list(values.get("tags"), max_items=20, max_length=30),
  534. "priority": _bounded_int(
  535. values.get("priority", 0), "知识优先级", minimum=-1000, maximum=1000
  536. ),
  537. "enabled": bool(values.get("enabled", True)),
  538. "source_candidate_id": clean_text(
  539. values.get("source_candidate_id"), max_length=64
  540. ),
  541. "source_references": values.get("source_references")
  542. if isinstance(values.get("source_references"), list)
  543. else [],
  544. "created_at": now,
  545. "updated_at": now,
  546. }
  547. await knowledgedb.insert_one(entry)
  548. return entry
  549. async def update_knowledge_entry(entry_id: str, values: dict[str, Any]) -> dict[str, Any]:
  550. filters = _scope(entry_id=str(entry_id))
  551. current = await knowledgedb.find_one(filters)
  552. if not current:
  553. raise AssistantDataError("knowledge_not_found", "未找到知识条目。")
  554. update: dict[str, Any] = {}
  555. if "question" in values:
  556. update["question"] = clean_text(values.get("question"), max_length=300, required=True)
  557. if "aliases" in values:
  558. update["aliases"] = _normalize_string_list(
  559. values.get("aliases"), max_items=20, max_length=200
  560. )
  561. if "keywords" in values:
  562. update["keywords"] = _normalize_string_list(
  563. values.get("keywords"), max_items=30, max_length=50
  564. )
  565. if "answer" in values:
  566. update["answer"] = clean_text(
  567. values.get("answer"), max_length=4000, required=True, preserve_lines=True
  568. )
  569. if "tags" in values:
  570. update["tags"] = _normalize_string_list(values.get("tags"), max_items=20, max_length=30)
  571. if "priority" in values:
  572. update["priority"] = _bounded_int(
  573. values.get("priority"), "知识优先级", minimum=-1000, maximum=1000
  574. )
  575. if "enabled" in values:
  576. update["enabled"] = bool(values.get("enabled"))
  577. if "source_candidate_id" in values:
  578. update["source_candidate_id"] = clean_text(
  579. values.get("source_candidate_id"), max_length=64
  580. )
  581. if "source_references" in values:
  582. references = values.get("source_references")
  583. if not isinstance(references, list):
  584. raise AssistantDataError("invalid_knowledge", "知识来源引用必须是数组。")
  585. update["source_references"] = references[:20]
  586. if not update:
  587. raise AssistantDataError("unchanged", "没有可保存的知识条目字段。")
  588. update["updated_at"] = utc_now()
  589. await knowledgedb.update_one(filters, {"$set": update})
  590. return await knowledgedb.find_one(filters) or current
  591. async def delete_knowledge_entry(entry_id: str) -> None:
  592. result = await knowledgedb.delete_one(_scope(entry_id=str(entry_id)))
  593. if not result.deleted_count:
  594. raise AssistantDataError("knowledge_not_found", "未找到知识条目。")
  595. async def list_knowledge_entries(
  596. connection_id: str,
  597. *,
  598. query: str = "",
  599. page: int = 1,
  600. page_size: int = 20,
  601. ) -> tuple[list[dict[str, Any]], int]:
  602. await ensure_assistant_indexes()
  603. filters: dict[str, Any] = _scope(connection_id=str(connection_id))
  604. normalized_query = clean_text(query, max_length=100)
  605. if normalized_query:
  606. pattern = re.escape(normalized_query)
  607. filters["$or"] = [
  608. {"question": {"$regex": pattern, "$options": "i"}},
  609. {"aliases": {"$regex": pattern, "$options": "i"}},
  610. {"keywords": {"$regex": pattern, "$options": "i"}},
  611. {"tags": {"$regex": pattern, "$options": "i"}},
  612. ]
  613. page = max(1, int(page))
  614. page_size = max(1, min(int(page_size), 100))
  615. total = await knowledgedb.count_documents(filters)
  616. cursor = (
  617. knowledgedb.find(filters)
  618. .sort([("priority", DESCENDING), ("updated_at", DESCENDING)])
  619. .skip((page - 1) * page_size)
  620. .limit(page_size)
  621. )
  622. return [item async for item in cursor], total
  623. def _knowledge_source_key(connection_id: str, source_type: str, chat_id: int) -> str:
  624. return f"{connection_id}:{source_type}:{int(chat_id)}"
  625. def normalize_knowledge_source(
  626. values: dict[str, Any], *, previous: dict[str, Any] | None = None
  627. ) -> dict[str, Any]:
  628. current = {**(previous or {}), **values}
  629. source_type = str(current.get("source_type") or "")
  630. if source_type not in VALID_SOURCE_TYPES:
  631. raise AssistantDataError("invalid_source", "知识来源类型无效。")
  632. try:
  633. chat_id = int(current.get("chat_id") or 0)
  634. except (TypeError, ValueError) as exc:
  635. raise AssistantDataError("invalid_source", "来源聊天 ID 必须是整数。") from exc
  636. if source_type == "business":
  637. chat_id = 0
  638. elif not chat_id:
  639. raise AssistantDataError("invalid_source", "频道或群组来源必须填写聊天 ID。")
  640. publication_mode = str(current.get("publication_mode") or "review")
  641. if publication_mode not in VALID_PUBLICATION_MODES:
  642. raise AssistantDataError("invalid_source", "知识发布方式无效。")
  643. author_policy = str(current.get("author_policy") or "admins_or_allowlist")
  644. if author_policy not in VALID_AUTHOR_POLICIES:
  645. raise AssistantDataError("invalid_source", "来源作者策略无效。")
  646. linked_chat_id = 0
  647. try:
  648. linked_chat_id = int(current.get("linked_chat_id") or 0)
  649. except (TypeError, ValueError) as exc:
  650. raise AssistantDataError("invalid_source", "关联讨论群 ID 必须是整数。") from exc
  651. return {
  652. "source_type": source_type,
  653. "chat_id": chat_id,
  654. "title": clean_text(
  655. current.get("title") or ("Business 日常对话" if source_type == "business" else ""),
  656. max_length=200,
  657. required=True,
  658. ),
  659. "publication_mode": publication_mode,
  660. "author_policy": author_policy,
  661. "allowed_user_ids": _normalize_int_list(current.get("allowed_user_ids")),
  662. "include_linked_chat": bool(current.get("include_linked_chat", False))
  663. if source_type == "channel"
  664. else False,
  665. "linked_chat_id": linked_chat_id if source_type == "channel" else 0,
  666. "linked_chat_title": clean_text(
  667. current.get("linked_chat_title"), max_length=200
  668. )
  669. if source_type == "channel"
  670. else "",
  671. "backfill_days": _bounded_int(
  672. current.get("backfill_days", 30), "历史回溯天数", minimum=0, maximum=365
  673. ),
  674. "backfill_limit": _bounded_int(
  675. current.get("backfill_limit", 100), "历史回溯条数", minimum=1, maximum=500
  676. ),
  677. "enabled": bool(current.get("enabled", True)),
  678. "access_status": clean_text(current.get("access_status"), max_length=64),
  679. }
  680. async def create_knowledge_source(
  681. connection_id: str, values: dict[str, Any]
  682. ) -> dict[str, Any]:
  683. await ensure_assistant_indexes()
  684. if not await get_business_connection(connection_id):
  685. raise AssistantDataError("connection_not_found", "未找到 Business 连接。")
  686. normalized = normalize_knowledge_source(values)
  687. now = utc_now()
  688. source = {
  689. "bot_id": BOT_PROFILE_ID,
  690. "source_id": uuid4().hex,
  691. "source_key": _knowledge_source_key(
  692. str(connection_id), normalized["source_type"], normalized["chat_id"]
  693. ),
  694. "connection_id": str(connection_id),
  695. **normalized,
  696. "sync_status": "idle",
  697. "last_error": "",
  698. "created_at": now,
  699. "updated_at": now,
  700. }
  701. try:
  702. await knowledge_sourcesdb.insert_one(source)
  703. except DuplicateKeyError as exc:
  704. raise AssistantDataError("source_exists", "该连接已配置相同的知识来源。") from exc
  705. return source
  706. async def get_knowledge_source(source_id: str) -> dict[str, Any] | None:
  707. await ensure_assistant_indexes()
  708. return await knowledge_sourcesdb.find_one(_scope(source_id=str(source_id)))
  709. async def update_knowledge_source(
  710. source_id: str, values: dict[str, Any]
  711. ) -> dict[str, Any]:
  712. current = await get_knowledge_source(source_id)
  713. if not current:
  714. raise AssistantDataError("source_not_found", "未找到知识来源。")
  715. for field in ("connection_id", "source_type", "chat_id"):
  716. if field in values and str(values[field]) != str(current[field]):
  717. raise AssistantDataError("immutable_source", "来源连接、类型和聊天 ID 不可修改。")
  718. normalized = normalize_knowledge_source(values, previous=current)
  719. await knowledge_sourcesdb.update_one(
  720. _scope(source_id=str(source_id)),
  721. {"$set": {**normalized, "updated_at": utc_now()}},
  722. )
  723. return await get_knowledge_source(source_id) or current
  724. async def list_knowledge_sources(
  725. connection_id: str,
  726. *,
  727. page: int = 1,
  728. page_size: int = 20,
  729. ) -> tuple[list[dict[str, Any]], int]:
  730. await ensure_assistant_indexes()
  731. filters = _scope(connection_id=str(connection_id))
  732. page = max(1, int(page))
  733. page_size = max(1, min(int(page_size), 100))
  734. total = await knowledge_sourcesdb.count_documents(filters)
  735. cursor = (
  736. knowledge_sourcesdb.find(filters)
  737. .sort("updated_at", DESCENDING)
  738. .skip((page - 1) * page_size)
  739. .limit(page_size)
  740. )
  741. return [item async for item in cursor], total
  742. async def list_sources_for_chat(
  743. chat_id: int, *, enabled_only: bool = True
  744. ) -> list[dict[str, Any]]:
  745. await ensure_assistant_indexes()
  746. values: dict[str, Any] = {
  747. "$or": [
  748. {"chat_id": int(chat_id)},
  749. {"include_linked_chat": True, "linked_chat_id": int(chat_id)},
  750. ]
  751. }
  752. if enabled_only:
  753. values["enabled"] = True
  754. filters = _scope(values)
  755. return [item async for item in knowledge_sourcesdb.find(filters)]
  756. async def list_enabled_sources_for_chat(chat_id: int) -> list[dict[str, Any]]:
  757. return await list_sources_for_chat(chat_id, enabled_only=True)
  758. async def list_enabled_business_sources(connection_id: str) -> list[dict[str, Any]]:
  759. await ensure_assistant_indexes()
  760. return [
  761. item
  762. async for item in knowledge_sourcesdb.find(
  763. _scope(
  764. connection_id=str(connection_id),
  765. source_type="business",
  766. enabled=True,
  767. )
  768. )
  769. ]
  770. async def list_business_sources(connection_id: str) -> list[dict[str, Any]]:
  771. await ensure_assistant_indexes()
  772. return [
  773. item
  774. async for item in knowledge_sourcesdb.find(
  775. _scope(connection_id=str(connection_id), source_type="business")
  776. )
  777. ]
  778. async def update_knowledge_source_sync(
  779. source_id: str,
  780. *,
  781. status: str,
  782. error: str = "",
  783. processed: int | None = None,
  784. ) -> dict[str, Any]:
  785. values: dict[str, Any] = {
  786. "sync_status": clean_text(status, max_length=32, required=True),
  787. "last_error": clean_text(error, max_length=1000),
  788. "updated_at": utc_now(),
  789. }
  790. if status == "running":
  791. values["last_sync_started_at"] = utc_now()
  792. if status in {"completed", "failed"}:
  793. values["last_sync_finished_at"] = utc_now()
  794. if processed is not None:
  795. values["last_sync_processed"] = max(0, int(processed))
  796. result = await knowledge_sourcesdb.update_one(
  797. _scope(source_id=str(source_id)), {"$set": values}
  798. )
  799. if not result.matched_count:
  800. raise AssistantDataError("source_not_found", "未找到知识来源。")
  801. return await get_knowledge_source(source_id) or {}
  802. async def touch_knowledge_source(source_id: str, *, error: str = "") -> None:
  803. await knowledge_sourcesdb.update_one(
  804. _scope(source_id=str(source_id)),
  805. {
  806. "$set": {
  807. "last_event_at": utc_now(),
  808. "last_error": clean_text(error, max_length=1000),
  809. "updated_at": utc_now(),
  810. }
  811. },
  812. )
  813. async def _disable_candidate_entries(filters: dict[str, Any]) -> None:
  814. entry_ids = [
  815. str(item.get("knowledge_entry_id") or "")
  816. async for item in knowledge_candidatesdb.find(filters)
  817. if item.get("knowledge_entry_id")
  818. ]
  819. if entry_ids:
  820. await knowledgedb.update_many(
  821. _scope({"entry_id": {"$in": entry_ids}}),
  822. {"$set": {"enabled": False, "updated_at": utc_now()}},
  823. )
  824. async def delete_knowledge_source(source_id: str) -> None:
  825. source = await get_knowledge_source(source_id)
  826. if not source:
  827. raise AssistantDataError("source_not_found", "未找到知识来源。")
  828. candidate_filters = _scope(source_id=str(source_id))
  829. await _disable_candidate_entries(candidate_filters)
  830. await knowledge_candidatesdb.update_many(
  831. candidate_filters,
  832. {"$set": {"status": "stale", "stale_reason": "source_deleted", "updated_at": utc_now()}},
  833. )
  834. await source_eventsdb.update_many(
  835. _scope(source_id=str(source_id)),
  836. {"$set": {"deleted": True, "delete_reason": "source_deleted", "updated_at": utc_now()}},
  837. )
  838. await knowledge_sourcesdb.delete_one(_scope(source_id=str(source_id)))
  839. def source_event_id_for(source_id: str, event_key: str) -> str:
  840. digest = hashlib.sha256(
  841. f"{BOT_PROFILE_ID}:{source_id}:{event_key}".encode()
  842. ).hexdigest()
  843. return digest[:32]
  844. async def record_source_event(
  845. source: dict[str, Any],
  846. *,
  847. event_key: str,
  848. content: str,
  849. metadata: dict[str, Any] | None = None,
  850. ) -> tuple[dict[str, Any], bool]:
  851. await ensure_assistant_indexes()
  852. cleaned_content = clean_text(
  853. content, max_length=20000, required=True, preserve_lines=True
  854. )
  855. cleaned_key = clean_text(event_key, max_length=300, required=True)
  856. event_metadata = metadata if isinstance(metadata, dict) else {}
  857. content_hash = hashlib.sha256(
  858. json.dumps(
  859. {"content": cleaned_content, "metadata": event_metadata},
  860. ensure_ascii=False,
  861. sort_keys=True,
  862. default=str,
  863. ).encode()
  864. ).hexdigest()
  865. filters = _scope(source_id=str(source["source_id"]), event_key=cleaned_key)
  866. current = await source_eventsdb.find_one(filters)
  867. if current and current.get("content_hash") == content_hash and not current.get("deleted"):
  868. return current, False
  869. now = utc_now()
  870. event_id = str(current.get("event_id")) if current else source_event_id_for(
  871. str(source["source_id"]), cleaned_key
  872. )
  873. version = int(current.get("version") or 0) + 1 if current else 1
  874. document = {
  875. "bot_id": BOT_PROFILE_ID,
  876. "event_id": event_id,
  877. "source_id": str(source["source_id"]),
  878. "connection_id": str(source["connection_id"]),
  879. "event_key": cleaned_key,
  880. "content": cleaned_content,
  881. "content_hash": content_hash,
  882. "metadata": event_metadata,
  883. "version": version,
  884. "deleted": False,
  885. "extraction_status": "pending",
  886. "extraction_error": "",
  887. "updated_at": now,
  888. "expires_at": now + timedelta(days=30),
  889. }
  890. await source_eventsdb.update_one(
  891. filters,
  892. {"$set": document, "$setOnInsert": {"created_at": now}},
  893. upsert=True,
  894. )
  895. return await source_eventsdb.find_one(filters) or document, True
  896. async def get_source_event(event_id: str) -> dict[str, Any] | None:
  897. await ensure_assistant_indexes()
  898. return await source_eventsdb.find_one(_scope(event_id=str(event_id)))
  899. async def get_source_event_by_key(
  900. source_id: str, event_key: str
  901. ) -> dict[str, Any] | None:
  902. await ensure_assistant_indexes()
  903. return await source_eventsdb.find_one(
  904. _scope(source_id=str(source_id), event_key=str(event_key))
  905. )
  906. async def update_source_event_extraction(
  907. event_id: str, *, status: str, error: str = "", item_count: int = 0
  908. ) -> None:
  909. await source_eventsdb.update_one(
  910. _scope(event_id=str(event_id)),
  911. {
  912. "$set": {
  913. "extraction_status": clean_text(status, max_length=32, required=True),
  914. "extraction_error": clean_text(error, max_length=1000),
  915. "extracted_item_count": max(0, int(item_count)),
  916. "extracted_at": utc_now(),
  917. "updated_at": utc_now(),
  918. }
  919. },
  920. )
  921. async def invalidate_source_event_candidates(
  922. event_id: str, *, reason: str
  923. ) -> None:
  924. filters = _scope(event_id=str(event_id))
  925. await _disable_candidate_entries(filters)
  926. await knowledge_candidatesdb.update_many(
  927. filters,
  928. {
  929. "$set": {
  930. "status": "stale",
  931. "stale_reason": clean_text(reason, max_length=100),
  932. "updated_at": utc_now(),
  933. }
  934. },
  935. )
  936. async def mark_source_event_deleted(source_id: str, event_key: str) -> bool:
  937. event = await source_eventsdb.find_one(
  938. _scope(source_id=str(source_id), event_key=str(event_key))
  939. )
  940. if not event:
  941. return False
  942. await source_eventsdb.update_one(
  943. _scope(event_id=str(event["event_id"])),
  944. {"$set": {"deleted": True, "updated_at": utc_now()}},
  945. )
  946. await invalidate_source_event_candidates(
  947. str(event["event_id"]), reason="source_message_deleted"
  948. )
  949. return True
  950. def _normalize_candidate_item(item: dict[str, Any]) -> dict[str, Any]:
  951. try:
  952. confidence = float(item.get("confidence") or 0)
  953. except (TypeError, ValueError):
  954. confidence = 0
  955. return {
  956. "question": clean_text(item.get("question"), max_length=300, required=True),
  957. "aliases": _normalize_string_list(item.get("aliases"), max_items=20, max_length=200),
  958. "keywords": _normalize_string_list(item.get("keywords"), max_items=30, max_length=50),
  959. "answer": clean_text(
  960. item.get("answer"), max_length=4000, required=True, preserve_lines=True
  961. ),
  962. "tags": _normalize_string_list(item.get("tags"), max_items=20, max_length=30),
  963. "confidence": max(0.0, min(confidence, 1.0)),
  964. }
  965. async def replace_event_candidates(
  966. source: dict[str, Any],
  967. event: dict[str, Any],
  968. items: list[dict[str, Any]],
  969. *,
  970. auto_publish: bool,
  971. ) -> list[dict[str, Any]]:
  972. await invalidate_source_event_candidates(
  973. str(event["event_id"]), reason="source_message_changed"
  974. )
  975. existing = {
  976. int(item.get("item_index") or 0): item
  977. async for item in knowledge_candidatesdb.find(
  978. _scope(event_id=str(event["event_id"]))
  979. )
  980. }
  981. results: list[dict[str, Any]] = []
  982. metadata = event.get("metadata") if isinstance(event.get("metadata"), dict) else {}
  983. for index, raw_item in enumerate(items[:5]):
  984. normalized = _normalize_candidate_item(raw_item)
  985. current = existing.get(index) or {}
  986. now = utc_now()
  987. candidate_id = str(current.get("candidate_id") or uuid4().hex)
  988. candidate = {
  989. "bot_id": BOT_PROFILE_ID,
  990. "candidate_id": candidate_id,
  991. "connection_id": str(source["connection_id"]),
  992. "source_id": str(source["source_id"]),
  993. "event_id": str(event["event_id"]),
  994. "item_index": index,
  995. **normalized,
  996. "status": "pending",
  997. "stale_reason": "",
  998. "knowledge_entry_id": str(current.get("knowledge_entry_id") or ""),
  999. "source_snapshot": {
  1000. "title": str(source.get("title") or ""),
  1001. "source_type": str(source.get("source_type") or ""),
  1002. "chat_id": int(metadata.get("chat_id") or source.get("chat_id") or 0),
  1003. "message_id": int(metadata.get("message_id") or 0),
  1004. "author_id": int(metadata.get("author_id") or 0),
  1005. },
  1006. "updated_at": now,
  1007. }
  1008. await knowledge_candidatesdb.update_one(
  1009. _scope(event_id=str(event["event_id"]), item_index=index),
  1010. {"$set": candidate, "$setOnInsert": {"created_at": now}},
  1011. upsert=True,
  1012. )
  1013. if auto_publish:
  1014. candidate = await publish_knowledge_candidate(candidate_id)
  1015. else:
  1016. candidate = await get_knowledge_candidate(candidate_id) or candidate
  1017. results.append(candidate)
  1018. await update_source_event_extraction(
  1019. str(event["event_id"]), status="completed", item_count=len(results)
  1020. )
  1021. return results
  1022. async def get_knowledge_candidate(candidate_id: str) -> dict[str, Any] | None:
  1023. await ensure_assistant_indexes()
  1024. return await knowledge_candidatesdb.find_one(
  1025. _scope(candidate_id=str(candidate_id))
  1026. )
  1027. async def publish_knowledge_candidate(candidate_id: str) -> dict[str, Any]:
  1028. candidate = await get_knowledge_candidate(candidate_id)
  1029. if not candidate:
  1030. raise AssistantDataError("candidate_not_found", "未找到知识候选。")
  1031. if candidate.get("status") == "stale":
  1032. raise AssistantDataError("candidate_stale", "来源已变更或删除,不能发布该候选。")
  1033. if not await get_knowledge_source(str(candidate.get("source_id") or "")):
  1034. await knowledge_candidatesdb.update_one(
  1035. _scope(candidate_id=str(candidate_id)),
  1036. {
  1037. "$set": {
  1038. "status": "stale",
  1039. "stale_reason": "source_deleted",
  1040. "updated_at": utc_now(),
  1041. }
  1042. },
  1043. )
  1044. raise AssistantDataError("source_not_found", "知识来源已删除,不能发布该候选。")
  1045. reference = {
  1046. **(candidate.get("source_snapshot") or {}),
  1047. "source_id": str(candidate.get("source_id") or ""),
  1048. "event_id": str(candidate.get("event_id") or ""),
  1049. }
  1050. values = {
  1051. key: candidate.get(key)
  1052. for key in ("question", "aliases", "keywords", "answer", "tags")
  1053. }
  1054. values.update(
  1055. {
  1056. "priority": round(float(candidate.get("confidence") or 0) * 100),
  1057. "enabled": True,
  1058. "source_candidate_id": str(candidate_id),
  1059. "source_references": [reference],
  1060. }
  1061. )
  1062. entry_id = str(candidate.get("knowledge_entry_id") or "")
  1063. if entry_id:
  1064. try:
  1065. entry = await update_knowledge_entry(entry_id, values)
  1066. except AssistantDataError as exc:
  1067. if exc.code != "knowledge_not_found":
  1068. raise
  1069. entry = await create_knowledge_entry(str(candidate["connection_id"]), values)
  1070. else:
  1071. entry = await create_knowledge_entry(str(candidate["connection_id"]), values)
  1072. await knowledge_candidatesdb.update_one(
  1073. _scope(candidate_id=str(candidate_id)),
  1074. {
  1075. "$set": {
  1076. "status": "published",
  1077. "knowledge_entry_id": str(entry["entry_id"]),
  1078. "published_at": utc_now(),
  1079. "updated_at": utc_now(),
  1080. }
  1081. },
  1082. )
  1083. return await get_knowledge_candidate(candidate_id) or candidate
  1084. async def reject_knowledge_candidate(candidate_id: str) -> dict[str, Any]:
  1085. candidate = await get_knowledge_candidate(candidate_id)
  1086. if not candidate:
  1087. raise AssistantDataError("candidate_not_found", "未找到知识候选。")
  1088. entry_id = str(candidate.get("knowledge_entry_id") or "")
  1089. if entry_id:
  1090. await knowledgedb.update_one(
  1091. _scope(entry_id=entry_id),
  1092. {"$set": {"enabled": False, "updated_at": utc_now()}},
  1093. )
  1094. await knowledge_candidatesdb.update_one(
  1095. _scope(candidate_id=str(candidate_id)),
  1096. {
  1097. "$set": {
  1098. "status": "rejected",
  1099. "rejected_at": utc_now(),
  1100. "updated_at": utc_now(),
  1101. }
  1102. },
  1103. )
  1104. return await get_knowledge_candidate(candidate_id) or candidate
  1105. async def list_knowledge_candidates(
  1106. connection_id: str,
  1107. *,
  1108. status: str = "",
  1109. source_id: str = "",
  1110. query: str = "",
  1111. page: int = 1,
  1112. page_size: int = 20,
  1113. ) -> tuple[list[dict[str, Any]], int]:
  1114. await ensure_assistant_indexes()
  1115. filters: dict[str, Any] = _scope(connection_id=str(connection_id))
  1116. if status:
  1117. if status not in VALID_CANDIDATE_STATUSES:
  1118. raise AssistantDataError("invalid_status", "知识候选状态无效。")
  1119. filters["status"] = status
  1120. if source_id:
  1121. filters["source_id"] = str(source_id)
  1122. normalized_query = clean_text(query, max_length=100)
  1123. if normalized_query:
  1124. pattern = re.escape(normalized_query)
  1125. filters["$or"] = [
  1126. {"question": {"$regex": pattern, "$options": "i"}},
  1127. {"answer": {"$regex": pattern, "$options": "i"}},
  1128. {"keywords": {"$regex": pattern, "$options": "i"}},
  1129. ]
  1130. page = max(1, int(page))
  1131. page_size = max(1, min(int(page_size), 100))
  1132. total = await knowledge_candidatesdb.count_documents(filters)
  1133. cursor = (
  1134. knowledge_candidatesdb.find(filters)
  1135. .sort("updated_at", DESCENDING)
  1136. .skip((page - 1) * page_size)
  1137. .limit(page_size)
  1138. )
  1139. return [item async for item in cursor], total
  1140. def _search_text(value: str) -> str:
  1141. return re.sub(r"[^\w\u3400-\u9fff]+", "", value.casefold())
  1142. async def match_knowledge(
  1143. connection_id: str, query: str, *, limit: int = 8
  1144. ) -> list[dict[str, Any]]:
  1145. await ensure_assistant_indexes()
  1146. normalized_query = _search_text(query)
  1147. if not normalized_query:
  1148. return []
  1149. entries = [
  1150. item
  1151. async for item in knowledgedb.find(
  1152. _scope(connection_id=str(connection_id), enabled=True)
  1153. )
  1154. ]
  1155. scored: list[tuple[int, dict[str, Any]]] = []
  1156. for entry in entries:
  1157. score = int(entry.get("priority") or 0)
  1158. phrases = [entry.get("question", ""), *(entry.get("aliases") or [])]
  1159. for phrase in phrases:
  1160. normalized_phrase = _search_text(str(phrase))
  1161. if not normalized_phrase:
  1162. continue
  1163. if normalized_query == normalized_phrase:
  1164. score += 1000
  1165. elif normalized_phrase in normalized_query:
  1166. score += 300 + min(len(normalized_phrase), 100)
  1167. elif normalized_query in normalized_phrase and len(normalized_query) >= 4:
  1168. score += 120
  1169. for keyword in entry.get("keywords") or []:
  1170. normalized_keyword = _search_text(str(keyword))
  1171. if normalized_keyword and normalized_keyword in normalized_query:
  1172. score += 80
  1173. if score > int(entry.get("priority") or 0):
  1174. scored.append((score, entry))
  1175. scored.sort(key=lambda item: (item[0], item[1].get("priority", 0)), reverse=True)
  1176. return [{**entry, "match_score": score} for score, entry in scored[: max(1, limit)]]
  1177. def classify_handoff(text: str) -> str | None:
  1178. normalized = text.casefold()
  1179. if any(term in normalized for term in HANDOFF_TERMS):
  1180. return "customer_requested_human"
  1181. if any(term in normalized for term in SENSITIVE_TERMS):
  1182. return "sensitive_request"
  1183. return None
  1184. def conversation_id_for(connection_id: str, chat_id: int) -> str:
  1185. digest = hashlib.sha256(f"{BOT_PROFILE_ID}:{connection_id}:{int(chat_id)}".encode()).hexdigest()
  1186. return digest[:32]
  1187. async def get_or_create_conversation(
  1188. connection_id: str,
  1189. chat_id: int,
  1190. *,
  1191. customer: dict[str, Any] | None = None,
  1192. ) -> dict[str, Any]:
  1193. await ensure_assistant_indexes()
  1194. conversation_id = conversation_id_for(connection_id, chat_id)
  1195. now = utc_now()
  1196. update: dict[str, Any] = {"updated_at": now, "last_message_at": now}
  1197. if customer:
  1198. update["customer"] = _public_user(customer)
  1199. await conversationsdb.update_one(
  1200. _scope(conversation_id=conversation_id),
  1201. {
  1202. "$set": update,
  1203. "$setOnInsert": {
  1204. "bot_id": BOT_PROFILE_ID,
  1205. "conversation_id": conversation_id,
  1206. "connection_id": str(connection_id),
  1207. "chat_id": int(chat_id),
  1208. "status": "auto",
  1209. "summary": "",
  1210. "handoff_reason": "",
  1211. "created_at": now,
  1212. },
  1213. },
  1214. upsert=True,
  1215. )
  1216. return await conversationsdb.find_one(_scope(conversation_id=conversation_id)) or {}
  1217. async def get_conversation(conversation_id: str) -> dict[str, Any] | None:
  1218. await ensure_assistant_indexes()
  1219. return await conversationsdb.find_one(
  1220. _scope(conversation_id=str(conversation_id))
  1221. )
  1222. async def get_conversation_by_chat(
  1223. connection_id: str, chat_id: int
  1224. ) -> dict[str, Any] | None:
  1225. await ensure_assistant_indexes()
  1226. return await conversationsdb.find_one(
  1227. _scope(connection_id=str(connection_id), chat_id=int(chat_id))
  1228. )
  1229. async def append_conversation_message(
  1230. conversation_id: str,
  1231. *,
  1232. direction: str,
  1233. telegram_message_id: int,
  1234. text: str,
  1235. sender_id: int = 0,
  1236. metadata: dict[str, Any] | None = None,
  1237. ) -> dict[str, Any]:
  1238. await ensure_assistant_indexes()
  1239. if direction not in {"incoming", "assistant", "human"}:
  1240. raise AssistantDataError("invalid_direction", "会话消息方向无效。")
  1241. now = utc_now()
  1242. message_key = f"{conversation_id}:{direction}:{int(telegram_message_id)}"
  1243. document = {
  1244. "bot_id": BOT_PROFILE_ID,
  1245. "message_key": message_key,
  1246. "conversation_id": str(conversation_id),
  1247. "direction": direction,
  1248. "telegram_message_id": int(telegram_message_id),
  1249. "sender_id": int(sender_id or 0),
  1250. "text": clean_text(text, max_length=12000, preserve_lines=True),
  1251. "metadata": metadata or {},
  1252. "created_at": now,
  1253. "expires_at": now + timedelta(days=30),
  1254. }
  1255. try:
  1256. await messagesdb.insert_one(document)
  1257. except DuplicateKeyError:
  1258. return await messagesdb.find_one(_scope(message_key=message_key)) or document
  1259. await conversationsdb.update_one(
  1260. _scope(conversation_id=str(conversation_id)),
  1261. {"$set": {"last_message_at": now, "updated_at": now}},
  1262. )
  1263. return document
  1264. async def claim_feishu_message_notification(
  1265. conversation_id: str, telegram_message_id: int, *, direction: str = "incoming"
  1266. ) -> bool:
  1267. if direction not in {"incoming", "human", "assistant"}:
  1268. raise AssistantDataError("invalid_direction", "会话消息方向无效。")
  1269. message_key = f"{conversation_id}:{direction}:{int(telegram_message_id)}"
  1270. stale_before = utc_now() - timedelta(minutes=2)
  1271. result = await messagesdb.find_one_and_update(
  1272. _scope(
  1273. {
  1274. "message_key": message_key,
  1275. "$or": [
  1276. {
  1277. "metadata.feishu_notification_status": {
  1278. "$nin": ["sending", "sent"]
  1279. }
  1280. },
  1281. {
  1282. "metadata.feishu_notification_status": "sending",
  1283. "metadata.feishu_notification_started_at": {
  1284. "$lte": stale_before
  1285. },
  1286. },
  1287. ],
  1288. }
  1289. ),
  1290. {
  1291. "$set": {
  1292. "metadata.feishu_notification_status": "sending",
  1293. "metadata.feishu_notification_started_at": utc_now(),
  1294. }
  1295. },
  1296. return_document=True,
  1297. )
  1298. return bool(result)
  1299. async def mark_feishu_message_notification(
  1300. conversation_id: str,
  1301. telegram_message_id: int,
  1302. *,
  1303. status: str,
  1304. error: str = "",
  1305. direction: str = "incoming",
  1306. ) -> None:
  1307. if status not in {"sent", "failed"}:
  1308. raise AssistantDataError("invalid_status", "飞书通知状态无效。")
  1309. if direction not in {"incoming", "human", "assistant"}:
  1310. raise AssistantDataError("invalid_direction", "会话消息方向无效。")
  1311. message_key = f"{conversation_id}:{direction}:{int(telegram_message_id)}"
  1312. await messagesdb.update_one(
  1313. _scope(message_key=message_key),
  1314. {
  1315. "$set": {
  1316. "metadata.feishu_notification_status": status,
  1317. "metadata.feishu_notification_error": clean_text(
  1318. error, max_length=1000
  1319. ),
  1320. "metadata.feishu_notification_finished_at": utc_now(),
  1321. }
  1322. },
  1323. )
  1324. async def recent_conversation_messages(
  1325. conversation_id: str, *, limit: int = 12
  1326. ) -> list[dict[str, Any]]:
  1327. cursor = (
  1328. messagesdb.find(_scope(conversation_id=str(conversation_id)))
  1329. .sort([("created_at", DESCENDING), ("_id", DESCENDING)])
  1330. .limit(max(1, min(int(limit), 50)))
  1331. )
  1332. items = [item async for item in cursor]
  1333. items.reverse()
  1334. return items
  1335. async def set_conversation_handoff(
  1336. conversation_id: str,
  1337. reason: str,
  1338. *,
  1339. summary: str = "",
  1340. ) -> dict[str, Any]:
  1341. now = utc_now()
  1342. update: dict[str, Any] = {
  1343. "status": "handoff",
  1344. "handoff_reason": clean_text(reason, max_length=200, required=True),
  1345. "paused_until": None,
  1346. "handoff_at": now,
  1347. "updated_at": now,
  1348. }
  1349. if summary:
  1350. update["summary"] = clean_text(summary, max_length=4000, preserve_lines=True)
  1351. result = await conversationsdb.find_one_and_update(
  1352. _scope(conversation_id=str(conversation_id)),
  1353. {"$set": update},
  1354. return_document=True,
  1355. )
  1356. if not result:
  1357. raise AssistantDataError("conversation_not_found", "未找到客户会话。")
  1358. return result
  1359. async def pause_conversation_for_human(
  1360. conversation_id: str, *, hours: int
  1361. ) -> dict[str, Any]:
  1362. now = utc_now()
  1363. result = await conversationsdb.find_one_and_update(
  1364. _scope(conversation_id=str(conversation_id)),
  1365. {
  1366. "$set": {
  1367. "status": "human_paused",
  1368. "handoff_reason": "human_reply_detected",
  1369. "paused_until": now + timedelta(hours=max(1, min(int(hours), 168))),
  1370. "updated_at": now,
  1371. }
  1372. },
  1373. return_document=True,
  1374. )
  1375. if not result:
  1376. raise AssistantDataError("conversation_not_found", "未找到客户会话。")
  1377. return result
  1378. async def resume_conversation(conversation_id: str) -> dict[str, Any]:
  1379. result = await conversationsdb.find_one_and_update(
  1380. _scope(conversation_id=str(conversation_id)),
  1381. {
  1382. "$set": {
  1383. "status": "auto",
  1384. "handoff_reason": "",
  1385. "paused_until": None,
  1386. "updated_at": utc_now(),
  1387. }
  1388. },
  1389. return_document=True,
  1390. )
  1391. if not result:
  1392. raise AssistantDataError("conversation_not_found", "未找到客户会话。")
  1393. return result
  1394. async def pause_conversation(conversation_id: str) -> dict[str, Any]:
  1395. return await set_conversation_handoff(conversation_id, "admin_paused")
  1396. async def close_conversation(conversation_id: str) -> dict[str, Any]:
  1397. result = await conversationsdb.find_one_and_update(
  1398. _scope(conversation_id=str(conversation_id)),
  1399. {"$set": {"status": "closed", "closed_at": utc_now(), "updated_at": utc_now()}},
  1400. return_document=True,
  1401. )
  1402. if not result:
  1403. raise AssistantDataError("conversation_not_found", "未找到客户会话。")
  1404. return result
  1405. async def clear_conversation(conversation_id: str) -> dict[str, Any]:
  1406. current = await get_conversation(conversation_id)
  1407. if not current:
  1408. raise AssistantDataError("conversation_not_found", "未找到客户会话。")
  1409. await messagesdb.delete_many(_scope(conversation_id=str(conversation_id)))
  1410. await conversationsdb.update_one(
  1411. _scope(conversation_id=str(conversation_id)),
  1412. {
  1413. "$set": {
  1414. "summary": "",
  1415. "handoff_reason": "",
  1416. "status": "auto",
  1417. "paused_until": None,
  1418. "cleared_at": utc_now(),
  1419. "updated_at": utc_now(),
  1420. }
  1421. },
  1422. )
  1423. return await get_conversation(conversation_id) or current
  1424. async def update_conversation_summary(
  1425. conversation_id: str, summary: str
  1426. ) -> None:
  1427. cleaned = clean_text(summary, max_length=4000, preserve_lines=True)
  1428. if cleaned:
  1429. await conversationsdb.update_one(
  1430. _scope(conversation_id=str(conversation_id)),
  1431. {"$set": {"summary": cleaned, "updated_at": utc_now()}},
  1432. )
  1433. async def list_conversations(
  1434. *,
  1435. connection_id: str = "",
  1436. status: str = "",
  1437. page: int = 1,
  1438. page_size: int = 20,
  1439. ) -> tuple[list[dict[str, Any]], int]:
  1440. await ensure_assistant_indexes()
  1441. filters: dict[str, Any] = _scope()
  1442. if connection_id:
  1443. filters["connection_id"] = str(connection_id)
  1444. if status:
  1445. if status not in {"auto", "handoff", "human_paused", "closed"}:
  1446. raise AssistantDataError("invalid_status", "会话状态无效。")
  1447. filters["status"] = status
  1448. page = max(1, int(page))
  1449. page_size = max(1, min(int(page_size), 100))
  1450. total = await conversationsdb.count_documents(filters)
  1451. cursor = (
  1452. conversationsdb.find(filters)
  1453. .sort("updated_at", DESCENDING)
  1454. .skip((page - 1) * page_size)
  1455. .limit(page_size)
  1456. )
  1457. return [item async for item in cursor], total
  1458. async def conversation_detail(conversation_id: str) -> dict[str, Any]:
  1459. conversation = await get_conversation(conversation_id)
  1460. if not conversation:
  1461. raise AssistantDataError("conversation_not_found", "未找到客户会话。")
  1462. conversation["messages"] = await recent_conversation_messages(
  1463. conversation_id, limit=50
  1464. )
  1465. return conversation
  1466. def _usage_day(now: datetime, timezone_name: str) -> str:
  1467. try:
  1468. zone = ZoneInfo(timezone_name)
  1469. except ZoneInfoNotFoundError:
  1470. zone = ZoneInfo("Asia/Shanghai")
  1471. return as_utc(now).astimezone(zone).date().isoformat()
  1472. async def reserve_ai_usage(
  1473. connection_id: str,
  1474. chat_id: int,
  1475. settings: dict[str, Any],
  1476. *,
  1477. now: datetime | None = None,
  1478. ) -> tuple[bool, str]:
  1479. await ensure_assistant_indexes()
  1480. current_time = as_utc(now or utc_now())
  1481. day = _usage_day(current_time, str(settings.get("timezone") or "Asia/Shanghai"))
  1482. account_limit = int(settings.get("account_daily_limit") or 200)
  1483. customer_limit = int(settings.get("customer_daily_limit") or 20)
  1484. account_filter = _scope(
  1485. connection_id=str(connection_id), chat_id=0, day=day
  1486. )
  1487. customer_filter = _scope(
  1488. connection_id=str(connection_id), chat_id=int(chat_id), day=day
  1489. )
  1490. for filters, usage_scope in (
  1491. (account_filter, "account"),
  1492. (customer_filter, "customer"),
  1493. ):
  1494. await usagedb.update_one(
  1495. filters,
  1496. {
  1497. "$setOnInsert": {
  1498. **filters,
  1499. "scope": usage_scope,
  1500. "call_count": 0,
  1501. "created_at": current_time,
  1502. "updated_at": current_time,
  1503. }
  1504. },
  1505. upsert=True,
  1506. )
  1507. account = await usagedb.find_one_and_update(
  1508. {**account_filter, "call_count": {"$lt": account_limit}},
  1509. {"$inc": {"call_count": 1}, "$set": {"updated_at": current_time}},
  1510. return_document=True,
  1511. )
  1512. if not account:
  1513. return False, "account_daily_limit"
  1514. customer = await usagedb.find_one_and_update(
  1515. {**customer_filter, "call_count": {"$lt": customer_limit}},
  1516. {"$inc": {"call_count": 1}, "$set": {"updated_at": current_time}},
  1517. return_document=True,
  1518. )
  1519. if not customer:
  1520. await usagedb.update_one(
  1521. {**account_filter, "call_count": {"$gt": 0}},
  1522. {"$inc": {"call_count": -1}, "$set": {"updated_at": current_time}},
  1523. )
  1524. return False, "customer_daily_limit"
  1525. return True, ""
  1526. async def usage_metrics(*, connection_id: str = "", day: str = "") -> dict[str, Any]:
  1527. await ensure_assistant_indexes()
  1528. match = _scope()
  1529. if connection_id:
  1530. match["connection_id"] = str(connection_id)
  1531. if day:
  1532. match["day"] = str(day)
  1533. calls = 0
  1534. customers: set[tuple[str, int]] = set()
  1535. customer_match = {
  1536. **match,
  1537. "scope": {"$ne": "account"},
  1538. "chat_id": {"$ne": 0},
  1539. "call_count": {"$gt": 0},
  1540. }
  1541. async for item in usagedb.find(customer_match):
  1542. calls += int(item.get("call_count") or 0)
  1543. customers.add((str(item.get("connection_id")), int(item.get("chat_id") or 0)))
  1544. conversation_filters = _scope()
  1545. if connection_id:
  1546. conversation_filters["connection_id"] = str(connection_id)
  1547. return {
  1548. "ai_calls": calls,
  1549. "customers": len(customers),
  1550. "conversations": await conversationsdb.count_documents(conversation_filters),
  1551. "handoffs": await conversationsdb.count_documents(
  1552. {**conversation_filters, "status": "handoff"}
  1553. ),
  1554. "human_paused": await conversationsdb.count_documents(
  1555. {**conversation_filters, "status": "human_paused"}
  1556. ),
  1557. }
  1558. async def claim_update(update_id: int) -> bool:
  1559. await ensure_assistant_indexes()
  1560. now = utc_now()
  1561. try:
  1562. await updatesdb.insert_one(
  1563. {
  1564. "bot_id": BOT_PROFILE_ID,
  1565. "update_id": int(update_id),
  1566. "status": "processing",
  1567. "attempts": 1,
  1568. "created_at": now,
  1569. "updated_at": now,
  1570. "expires_at": now + timedelta(days=7),
  1571. }
  1572. )
  1573. return True
  1574. except DuplicateKeyError:
  1575. filters = _scope(update_id=int(update_id))
  1576. current = await updatesdb.find_one(filters) or {}
  1577. if current.get("status") == "done":
  1578. return False
  1579. await updatesdb.update_one(
  1580. filters,
  1581. {"$inc": {"attempts": 1}, "$set": {"updated_at": now}},
  1582. )
  1583. return True
  1584. async def mark_update_done(update_id: int) -> None:
  1585. await updatesdb.update_one(
  1586. _scope(update_id=int(update_id)),
  1587. {"$set": {"status": "done", "updated_at": utc_now(), "last_error": ""}},
  1588. )
  1589. async def mark_update_failed(update_id: int, error: str) -> int:
  1590. await updatesdb.update_one(
  1591. _scope(update_id=int(update_id)),
  1592. {"$set": {"status": "failed", "updated_at": utc_now(), "last_error": error[:1000]}},
  1593. )
  1594. current = await updatesdb.find_one(_scope(update_id=int(update_id))) or {}
  1595. return int(current.get("attempts") or 1)
  1596. async def dead_letter_update(update: dict[str, Any], error: str) -> None:
  1597. update_id = int(update.get("update_id") or 0)
  1598. now = utc_now()
  1599. await deadlettersdb.update_one(
  1600. _scope(update_id=update_id),
  1601. {
  1602. "$set": {
  1603. "bot_id": BOT_PROFILE_ID,
  1604. "payload": update,
  1605. "error": error[:2000],
  1606. "updated_at": now,
  1607. "expires_at": now + timedelta(days=30),
  1608. },
  1609. "$setOnInsert": {"created_at": now},
  1610. },
  1611. upsert=True,
  1612. )
  1613. await mark_update_done(update_id)
  1614. async def runtime_status() -> dict[str, Any]:
  1615. await ensure_assistant_indexes()
  1616. stored = await runtimedb.find_one(
  1617. _scope(runtime_id="business_assistant")
  1618. ) or {}
  1619. return {
  1620. "runtime_id": "business_assistant",
  1621. "polling_state": "stopped",
  1622. "business_mode_supported": False,
  1623. "webhook_conflict": False,
  1624. "webhook_url": "",
  1625. "model_configured": False,
  1626. "last_error": "",
  1627. **{key: value for key, value in stored.items() if key != "_id"},
  1628. }
  1629. async def update_runtime_status(values: dict[str, Any]) -> dict[str, Any]:
  1630. await ensure_assistant_indexes()
  1631. insert_values = {
  1632. "bot_id": BOT_PROFILE_ID,
  1633. "runtime_id": "business_assistant",
  1634. "created_at": utc_now(),
  1635. }
  1636. if "offset" not in values:
  1637. insert_values["offset"] = 0
  1638. await runtimedb.update_one(
  1639. _scope(runtime_id="business_assistant"),
  1640. {
  1641. "$set": {**values, "updated_at": utc_now()},
  1642. "$setOnInsert": insert_values,
  1643. },
  1644. upsert=True,
  1645. )
  1646. return await runtime_status()
  1647. async def load_update_offset() -> int:
  1648. return int((await runtime_status()).get("offset") or 0)
  1649. async def save_update_offset(offset: int) -> None:
  1650. await update_runtime_status({"offset": int(offset), "last_poll_at": utc_now()})
  1651. async def due_digest_connections(now: datetime | None = None) -> list[dict[str, Any]]:
  1652. current = as_utc(now or utc_now())
  1653. results: list[dict[str, Any]] = []
  1654. async for settings in settingsdb.find(
  1655. _scope(digest_enabled=True, assistant_enabled=True)
  1656. ):
  1657. try:
  1658. local = current.astimezone(ZoneInfo(str(settings.get("timezone"))))
  1659. except ZoneInfoNotFoundError:
  1660. local = current.astimezone(ZoneInfo("Asia/Shanghai"))
  1661. day = local.date().isoformat()
  1662. if local.strftime("%H:%M") != str(settings.get("digest_time") or "09:00"):
  1663. continue
  1664. if settings.get("last_digest_day") == day:
  1665. continue
  1666. connection = await get_business_connection(settings["connection_id"])
  1667. if connection and connection.get("is_enabled"):
  1668. results.append({"connection": connection, "settings": settings, "day": day})
  1669. return results
  1670. async def mark_digest_sent(connection_id: str, day: str) -> None:
  1671. await settingsdb.update_one(
  1672. _scope(connection_id=str(connection_id)),
  1673. {"$set": {"last_digest_day": str(day), "last_digest_at": utc_now()}},
  1674. )