diff --git a/.gitignore b/.gitignore index db41c35..f393843 100644 --- a/.gitignore +++ b/.gitignore @@ -31,3 +31,4 @@ Thumbs.db # Compiled binaries and output find_all_keys_macos decoded_voices/ +voice_transcriptions.json diff --git a/mcp_server.py b/mcp_server.py index c464a06..15ec657 100644 --- a/mcp_server.py +++ b/mcp_server.py @@ -6,7 +6,7 @@ Runs on Windows Python (needs access to D:\ WeChat databases). """ import io -import os, sys, json, time, sqlite3, tempfile, struct, hashlib, atexit, re +import os, sys, json, time, sqlite3, tempfile, struct, hashlib, atexit, re, threading import wave import hmac as hmac_mod from contextlib import closing @@ -1890,9 +1890,92 @@ def decode_voice(chat_name: str, local_id: int) -> str: ) +# ============ 语音转录缓存 ============ +# +# Whisper 转录耗时(CPU 下每条数秒到数十秒),且结果是确定性的 +# (同一段 voice_data → 同一段 text),非常适合缓存。 +# +# 缓存 key 用 json.dumps([username, local_id]):local_id 在单个 username 下 +# 稳定唯一,套一层 JSON 序列化保证 username 里若含分隔符也不会与其它条目碰撞。 +# +# 写入走 temp + os.replace 原子替换,避免进程中途被杀导致整份缓存损坏 +# (Whisper 的单次代价远高于 DBCache,破档不可接受)。 +# +# 条目里记录 model_size:Whisper 升级默认模型后,旧条目自动视为失效并重跑。 + +VOICE_TRANSCRIPTION_CACHE_FILE = os.path.join(SCRIPT_DIR, "voice_transcriptions.json") + +_voice_transcription_cache = None # 懒加载 dict;None 表示尚未加载 +_voice_transcription_cache_lock = threading.Lock() +_voice_transcription_save_warned = False # 写失败仅首次写 stderr,避免刷屏 + + +def _voice_transcription_cache_key(username, local_id): + """构造缓存 key。用 json.dumps 兜底 username 里可能出现的分隔符。""" + return json.dumps([username, int(local_id)], ensure_ascii=False) + + +def _load_voice_transcription_cache(): + """加载缓存到模块级 dict,返回该 dict。 + + 文件不存在 → 空 dict。JSON 损坏或 payload 非 dict → 空 dict + (与上游 DBCache 的容错风格一致:缓存坏了不要拖垮工具调用)。 + """ + global _voice_transcription_cache + with _voice_transcription_cache_lock: + if _voice_transcription_cache is not None: + return _voice_transcription_cache + if not os.path.exists(VOICE_TRANSCRIPTION_CACHE_FILE): + _voice_transcription_cache = {} + return _voice_transcription_cache + try: + with open(VOICE_TRANSCRIPTION_CACHE_FILE, encoding="utf-8") as f: + loaded = json.load(f) + _voice_transcription_cache = loaded if isinstance(loaded, dict) else {} + except (json.JSONDecodeError, OSError): + _voice_transcription_cache = {} + return _voice_transcription_cache + + +def _save_voice_transcription_cache(): + """持久化缓存到磁盘。 + + - 原子写:先写 .tmp 再 os.replace,避免 crash 中途留下半截文件。 + - 未加载过也允许保存:此时把 module 状态初始化为空 dict,避免上层 + 代码因调用顺序错误而静默丢数据。 + - OSError 不抛:避免转录成功但落盘失败时让工具调用也失败;但首次 + 失败会在 stderr 打一行警告,用户知道磁盘满 / 权限问题需要处理。 + """ + global _voice_transcription_cache, _voice_transcription_save_warned + with _voice_transcription_cache_lock: + if _voice_transcription_cache is None: + _voice_transcription_cache = {} + tmp_path = VOICE_TRANSCRIPTION_CACHE_FILE + ".tmp" + try: + with open(tmp_path, "w", encoding="utf-8") as f: + json.dump(_voice_transcription_cache, f, ensure_ascii=False) + os.replace(tmp_path, VOICE_TRANSCRIPTION_CACHE_FILE) + except OSError as exc: + if not _voice_transcription_save_warned: + print( + f"[voice_cache] 写入失败(后续不再提示): {exc}", + file=sys.stderr, + flush=True, + ) + _voice_transcription_save_warned = True + # 清理可能残留的 .tmp + try: + if os.path.exists(tmp_path): + os.unlink(tmp_path) + except OSError: + pass + + +DEFAULT_WHISPER_MODEL = "base" + _whisper_model = None -def _get_whisper_model(model_size="base"): +def _get_whisper_model(model_size=DEFAULT_WHISPER_MODEL): global _whisper_model if _whisper_model is None: import whisper @@ -1904,8 +1987,10 @@ def _get_whisper_model(model_size="base"): def transcribe_voice(chat_name: str, local_id: int) -> str: """将微信语音消息转录为文字(自动检测语言,保留原语言)。 - 会先解码 SILK 语音为 WAV,再用 Whisper 转录。 - 首次运行会下载 Whisper 模型(约 145MB)。 + 首次转录会先解码 SILK 语音为 WAV,再用 Whisper 转录;结果缓存到 + voice_transcriptions.json,重复调用直接返回缓存(跳过 SILK 解码 + 和 Whisper 推理)。若 Whisper 默认模型升级(如 base → small), + 旧条目自动视为失效并重新转录。首次运行会下载 Whisper 模型(约 145MB)。 依赖: pip install silk-python openai-whisper (silk-python 的 import 名为 pysilk) @@ -1913,6 +1998,29 @@ def transcribe_voice(chat_name: str, local_id: int) -> str: chat_name: 聊天对象的名字、备注名或wxid local_id: 语音消息的 local_id(从 get_voice_messages 获取) """ + username = resolve_username(chat_name) + if not username: + return f"找不到聊天对象: {chat_name}" + + cache_key = _voice_transcription_cache_key(username, local_id) + cache = _load_voice_transcription_cache() + entry = cache.get(cache_key) + if ( + isinstance(entry, dict) + and "text" in entry + and entry.get("model_size") == DEFAULT_WHISPER_MODEL + ): + # 命中缓存:跳过 DB 查询、SILK 解码、Whisper 推理。 + # 条目里存了 create_time,即使源 DB 中消息已被清理仍能返回历史转录。 + lang = entry.get("language", "unknown") + cached_ts = entry.get("create_time") + if isinstance(cached_ts, int): + time_label = datetime.fromtimestamp(cached_ts).strftime('%Y-%m-%d %H:%M') + else: + time_label = "-" + return f"[{time_label}] ({lang})\n{entry['text']}" + + # 未命中:只有这条路径才需要 whisper / pysilk。 try: import whisper # noqa: F401 except ImportError: @@ -1922,10 +2030,6 @@ def transcribe_voice(chat_name: str, local_id: int) -> str: except ImportError: return "缺少依赖: pip install silk-python" - username = resolve_username(chat_name) - if not username: - return f"找不到聊天对象: {chat_name}" - row = _fetch_voice_row(username, local_id) if row is None: return f"找不到 local_id={local_id} 的语音消息" @@ -1938,6 +2042,16 @@ def transcribe_voice(chat_name: str, local_id: int) -> str: lang = result.get("language", "unknown") text = result.get("text", "").strip() + # 写缓存:即使 text 为空也缓存(Whisper 偶尔对静音/极短片段返回空), + # 配合 model_size 字段,升级模型后会自动重转,避免永久钉死空结果。 + cache[cache_key] = { + "text": text, + "language": lang, + "create_time": int(create_time), + "model_size": DEFAULT_WHISPER_MODEL, + } + _save_voice_transcription_cache() + time_label = datetime.fromtimestamp(create_time).strftime('%Y-%m-%d %H:%M') return f"[{time_label}] ({lang})\n{text}" diff --git a/tests/test_voice_transcription_cache.py b/tests/test_voice_transcription_cache.py new file mode 100644 index 0000000..957596b --- /dev/null +++ b/tests/test_voice_transcription_cache.py @@ -0,0 +1,274 @@ +import json +import os +import tempfile +import threading +import unittest +from unittest.mock import patch + +import mcp_server + + +class _CacheIsolationMixin: + """所有测试共享:隔离 module-level 缓存状态 + 指向 tempdir 的 cache 文件。""" + + def setUp(self): + self._saved_cache = mcp_server._voice_transcription_cache + self._saved_path = mcp_server.VOICE_TRANSCRIPTION_CACHE_FILE + self._saved_warned = mcp_server._voice_transcription_save_warned + + mcp_server._voice_transcription_cache = None + mcp_server._voice_transcription_save_warned = False + + self._tmp = tempfile.TemporaryDirectory() + self.addCleanup(self._tmp.cleanup) + mcp_server.VOICE_TRANSCRIPTION_CACHE_FILE = os.path.join( + self._tmp.name, "voice_transcriptions.json" + ) + + def tearDown(self): + mcp_server._voice_transcription_cache = self._saved_cache + mcp_server.VOICE_TRANSCRIPTION_CACHE_FILE = self._saved_path + mcp_server._voice_transcription_save_warned = self._saved_warned + + +class VoiceTranscriptionCachePersistenceTests(_CacheIsolationMixin, unittest.TestCase): + """_load_voice_transcription_cache / _save_voice_transcription_cache 的持久化行为。""" + + def test_load_missing_file_returns_empty_dict(self): + self.assertEqual(mcp_server._load_voice_transcription_cache(), {}) + + def test_save_and_reload_roundtrip(self): + cache = mcp_server._load_voice_transcription_cache() + cache["wxid_foo:42"] = { + "text": "你好", + "language": "zh", + "create_time": 1700000000, + "model_size": "base", + } + mcp_server._save_voice_transcription_cache() + + # 强制下一次 load 从磁盘读 + mcp_server._voice_transcription_cache = None + reloaded = mcp_server._load_voice_transcription_cache() + self.assertEqual(reloaded["wxid_foo:42"]["text"], "你好") + self.assertEqual(reloaded["wxid_foo:42"]["language"], "zh") + + def test_corrupt_file_returns_empty_dict(self): + with open(mcp_server.VOICE_TRANSCRIPTION_CACHE_FILE, "w", encoding="utf-8") as f: + f.write("{{ not valid json") + self.assertEqual(mcp_server._load_voice_transcription_cache(), {}) + + def test_non_dict_payload_returns_empty_dict(self): + with open(mcp_server.VOICE_TRANSCRIPTION_CACHE_FILE, "w", encoding="utf-8") as f: + json.dump(["not", "a", "dict"], f) + self.assertEqual(mcp_server._load_voice_transcription_cache(), {}) + + def test_utf8_preserved_on_disk(self): + # ensure_ascii=False 必须生效,否则中文会被转义成 \uXXXX + cache = mcp_server._load_voice_transcription_cache() + cache["wxid_bar:1"] = {"text": "中文测试", "language": "zh"} + mcp_server._save_voice_transcription_cache() + + with open(mcp_server.VOICE_TRANSCRIPTION_CACHE_FILE, "rb") as f: + raw = f.read() + self.assertIn("中文测试".encode("utf-8"), raw) + + def test_save_without_prior_load_persists_empty_dict(self): + # 从未 load 过就直接 save:应落盘一个空 dict,而不是静默丢弃。 + mcp_server._voice_transcription_cache = None + mcp_server._save_voice_transcription_cache() + self.assertTrue(os.path.exists(mcp_server.VOICE_TRANSCRIPTION_CACHE_FILE)) + with open(mcp_server.VOICE_TRANSCRIPTION_CACHE_FILE, encoding="utf-8") as f: + self.assertEqual(json.load(f), {}) + + +class VoiceTranscriptionCacheAtomicityTests(_CacheIsolationMixin, unittest.TestCase): + """原子写 + crash-during-save 行为。""" + + def test_write_is_atomic_via_rename(self): + # 先写入一份已有缓存 + cache = mcp_server._load_voice_transcription_cache() + cache["wxid_x:1"] = {"text": "initial", "language": "zh", "model_size": "base"} + mcp_server._save_voice_transcription_cache() + + # 模拟:写 .tmp 正常但 os.replace 阶段失败 + original_replace = os.replace + + def flaky_replace(src, dst): + raise OSError("disk full during rename") + + cache["wxid_x:1"] = {"text": "MUTATED", "language": "zh", "model_size": "base"} + with patch.object(os, "replace", side_effect=flaky_replace): + mcp_server._save_voice_transcription_cache() # 不应抛 + + # 磁盘上应仍然是 initial,不是 MUTATED,也不是损坏的半截文件 + with open(mcp_server.VOICE_TRANSCRIPTION_CACHE_FILE, encoding="utf-8") as f: + disk = json.load(f) + self.assertEqual(disk["wxid_x:1"]["text"], "initial") + + # .tmp 应该被清理,避免污染目录 + tmp_path = mcp_server.VOICE_TRANSCRIPTION_CACHE_FILE + ".tmp" + # 注:patch 生效期间 os.replace 失败,finally 里会尝试 unlink + _ = original_replace # 防 lint 警告 + self.assertFalse(os.path.exists(tmp_path)) + + def test_early_save_error_preserves_existing_file(self): + # json.dump 在 .tmp 上抛异常时(模拟磁盘满 / 权限问题),主文件应保持原样; + # 注意此测试不是"写到一半中断"而是"写前就失败"的场景。 + cache = mcp_server._load_voice_transcription_cache() + cache["wxid_y:1"] = {"text": "survives", "language": "zh", "model_size": "base"} + mcp_server._save_voice_transcription_cache() + + cache["wxid_y:1"] = {"text": "DO NOT SEE", "language": "zh", "model_size": "base"} + + def boom(*args, **kwargs): + raise OSError("disk full") + + with patch.object(mcp_server.json, "dump", side_effect=boom): + mcp_server._save_voice_transcription_cache() # 静默降级,不抛 + + # 主文件没被破坏:仍然可 json.load 出原先内容 + mcp_server._voice_transcription_cache = None + reloaded = mcp_server._load_voice_transcription_cache() + self.assertEqual(reloaded["wxid_y:1"]["text"], "survives") + + +class VoiceTranscriptionCacheConcurrencyTests(_CacheIsolationMixin, unittest.TestCase): + """多线程下的 load/save 行为。""" + + def test_concurrent_load_returns_same_dict_instance(self): + # 16 个线程同时触发首次 load,应当只实际化一份 dict(lock 生效) + barrier = threading.Barrier(16) + results = [] + results_lock = threading.Lock() + + def worker(): + barrier.wait() + d = mcp_server._load_voice_transcription_cache() + with results_lock: + results.append(id(d)) + + threads = [threading.Thread(target=worker) for _ in range(16)] + for t in threads: + t.start() + for t in threads: + t.join() + + self.assertEqual(len(set(results)), 1, "并发 load 应返回同一个 dict 对象") + + def test_concurrent_save_does_not_corrupt(self): + # 多个线程同时 save,磁盘上最终文件必须是合法 JSON(原子写 + lock 保障) + cache = mcp_server._load_voice_transcription_cache() + for i in range(100): + cache[f"wxid_z:{i}"] = { + "text": f"msg-{i}", + "language": "zh", + "model_size": "base", + } + + def worker(): + mcp_server._save_voice_transcription_cache() + + threads = [threading.Thread(target=worker) for _ in range(8)] + for t in threads: + t.start() + for t in threads: + t.join() + + with open(mcp_server.VOICE_TRANSCRIPTION_CACHE_FILE, encoding="utf-8") as f: + disk = json.load(f) # 必须能解析 + self.assertEqual(len(disk), 100) + + +class TranscribeVoiceCacheHitTests(_CacheIsolationMixin, unittest.TestCase): + """transcribe_voice 的缓存命中 / 失效路径。""" + + def _seed(self, key, entry): + cache = mcp_server._load_voice_transcription_cache() + cache[key] = entry + mcp_server._save_voice_transcription_cache() + + def test_cache_hit_skips_fetch_and_transcribe(self): + key = mcp_server._voice_transcription_cache_key("wxid_test", 7) + self._seed(key, { + "text": "缓存命中文本", + "language": "zh", + "create_time": 1700000000, + "model_size": mcp_server.DEFAULT_WHISPER_MODEL, + }) + + with patch.object(mcp_server, "resolve_username", return_value="wxid_test") as mock_resolve, \ + patch.object(mcp_server, "_fetch_voice_row") as mock_fetch, \ + patch.object(mcp_server, "_silk_to_wav") as mock_silk, \ + patch.object(mcp_server, "_get_whisper_model") as mock_model: + result = mcp_server.transcribe_voice("test_contact", 7) + + mock_resolve.assert_called_once_with("test_contact") + mock_fetch.assert_not_called() + mock_silk.assert_not_called() + mock_model.assert_not_called() + self.assertIn("缓存命中文本", result) + self.assertIn("(zh)", result) + + def test_cache_hit_uses_placeholder_when_create_time_missing(self): + # 旧条目若没有 create_time 字段,不应崩溃 + key = mcp_server._voice_transcription_cache_key("wxid_test", 8) + self._seed(key, { + "text": "历史条目", + "language": "zh", + "model_size": mcp_server.DEFAULT_WHISPER_MODEL, + }) + + with patch.object(mcp_server, "resolve_username", return_value="wxid_test"), \ + patch.object(mcp_server, "_fetch_voice_row") as mock_fetch: + result = mcp_server.transcribe_voice("test_contact", 8) + + mock_fetch.assert_not_called() + self.assertIn("历史条目", result) + + def test_cache_hit_returns_empty_text_without_retranscribing(self): + # Whisper 返回空也要缓存;再次调用应直接返回空,不进入 miss 路径 + key = mcp_server._voice_transcription_cache_key("wxid_test", 9) + self._seed(key, { + "text": "", + "language": "zh", + "create_time": 1700000000, + "model_size": mcp_server.DEFAULT_WHISPER_MODEL, + }) + + with patch.object(mcp_server, "resolve_username", return_value="wxid_test"), \ + patch.object(mcp_server, "_fetch_voice_row") as mock_fetch: + result = mcp_server.transcribe_voice("test_contact", 9) + + mock_fetch.assert_not_called() + self.assertIn("(zh)", result) + + def test_model_mismatch_is_treated_as_miss(self): + # 缓存条目的 model_size 和当前 DEFAULT_WHISPER_MODEL 不一致时, + # 不应命中;进入 miss 路径(这里无 whisper 依赖,应落到"缺少依赖"分支)。 + key = mcp_server._voice_transcription_cache_key("wxid_test", 10) + self._seed(key, { + "text": "旧模型结果", + "language": "zh", + "create_time": 1700000000, + "model_size": "OUTDATED_MODEL", + }) + + with patch.object(mcp_server, "resolve_username", return_value="wxid_test"), \ + patch.dict("sys.modules", {"whisper": None}): + # whisper=None 时 `import whisper` 触发 ImportError + result = mcp_server.transcribe_voice("test_contact", 10) + + # 走了 miss 路径 → 返回缺依赖提示,而不是返回旧缓存文本 + self.assertNotIn("旧模型结果", result) + self.assertIn("缺少依赖", result) + + def test_cache_key_handles_colon_in_username(self): + # 若上游未来的 resolve_username 放出带 ':' 的 username,也不会和其他条目冲突 + key_a = mcp_server._voice_transcription_cache_key("wxid:foo", 1) + key_b = mcp_server._voice_transcription_cache_key("wxid", 1) + self.assertNotEqual(key_a, key_b) + + +if __name__ == "__main__": + unittest.main()