Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion .github/workflows/ai-cross-platform.yml
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,11 @@ jobs:
- name: AI and local search regression
shell: bash
# 首次失败立即输出具体断言,避免超时取消后只留下进度标记。
run: uv run pytest -x -vv --tb=short --durations=15 tests/test_ai*.py tests/test_local_search*.py
# 由 Python 展开 glob,避免 Windows 默认 PowerShell 将通配符原样传给 pytest。
run: >-
uv run python -c
"import glob, pytest; raise SystemExit(pytest.main(['-x', '-vv', '--tb=short', '--durations=15',
*sorted(glob.glob('tests/test_ai*.py') + glob.glob('tests/test_local_search*.py'))]))"
- name: Chat calendar and date navigation regression
run: >-
uv run pytest -q
Expand Down
52 changes: 50 additions & 2 deletions src/wechat_decrypt_tool/ai/lifecycle.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,16 @@
"""AI 服务统一启停诊断,独立于其他后台服务的生命周期。"""
import logging
import threading
from typing import Any
from importlib.metadata import PackageNotFoundError, version

from .diagnostics import observed, event

STORE_MAINTENANCE_INTERVAL_SECONDS = 6 * 60 * 60
_maintenance_lock = threading.Lock()
_maintenance_stop: threading.Event | None = None
_maintenance_threads: list[threading.Thread] = []


def _maintain_store(store, name):
"""在后台回收过期事件并压缩数据库,避免大型遗留库拖慢启动。"""
Expand All @@ -16,6 +22,48 @@ def _maintain_store(store, name):
event('storage.maintenance.failed', level=logging.WARNING, component=name, error=error)


def _maintenance_loop(store: Any, name: str, stop: threading.Event,
interval: float = STORE_MAINTENANCE_INTERVAL_SECONDS) -> None:
"""定期维护存储;等待使用 stop,以便应用关闭时立即退出等待。"""
while not stop.is_set():
_maintain_store(store, name)
stop.wait(max(0.1, interval))


def _start_store_maintenance(stores: tuple[tuple[str, Any], ...]) -> None:
"""启动一组唯一的维护线程,避免重复触发多个生命周期钩子。"""
global _maintenance_stop, _maintenance_threads
with _maintenance_lock:
if _maintenance_stop is not None and any(thread.is_alive() for thread in _maintenance_threads):
return
_maintenance_stop = threading.Event()
_maintenance_threads = []
for name, store in stores:
thread = threading.Thread(
target=_maintenance_loop,
args=(store, name, _maintenance_stop),
name=f'ai-store-maintenance-{name}',
daemon=True,
)
_maintenance_threads.append(thread)
thread.start()


def _stop_store_maintenance() -> None:
"""通知维护线程停止,并等待短暂时间让正在执行的维护收尾。"""
global _maintenance_stop, _maintenance_threads
with _maintenance_lock:
stop = _maintenance_stop
threads = list(_maintenance_threads)
_maintenance_stop = None
_maintenance_threads = []
if stop is None:
return
stop.set()
for thread in threads:
thread.join(timeout=2)


@observed('lifecycle.start')
async def start_services():
for package in ('deepagents', 'langchain-openai', 'langchain-anthropic', 'langgraph', 'onnxruntime', 'sqlite-vec', 'tokenizers'):
Expand All @@ -30,15 +78,15 @@ async def start_services():
get_ai_service().start()
await get_agent_service().start()
await get_local_search().start()
for name, store in (('summary', get_ai_service().store), ('search', get_local_search().store)):
threading.Thread(target=_maintain_store, args=(store, name), name=f'ai-store-maintenance-{name}', daemon=True).start()
_start_store_maintenance((('summary', get_ai_service().store), ('search', get_local_search().store)))


@observed('lifecycle.stop')
async def stop_services():
from .service import get_ai_service
from .agent_service import get_agent_service
from ..local_search.service import get_local_search
_stop_store_maintenance()
for name, factory in (('search',get_local_search),('summary',get_ai_service),('agent',get_agent_service)):
try:
await factory().stop()
Expand Down
4 changes: 2 additions & 2 deletions tests/test_ai_message_pages.py
Original file line number Diff line number Diff line change
Expand Up @@ -371,8 +371,8 @@ def record(account, kind, body, unique_key=None, replace=False):
with service.index('a').connection() as db:
assert db.execute('SELECT count(*) FROM messages').fetchone()[0] == 1250
events = service.store.events()
assert any(e['kind'] == 'local_search_index' and e['body'].get('stage') == 'reading'
and e['body'].get('read_count', 0) > e['body']['processed'] for e in emitted)
assert any(e.get('stage') == 'reading'
and e.get('read_count', 0) > e.get('processed', 0) for e in emitted)
# 同一任务在事件表只保留最新一行,且不携带可重建的大字段。
stored = [e for e in events if e['kind'] == 'local_search_index' and e['body'].get('id') == job['id']]
assert len(stored) == 1
Expand Down
19 changes: 19 additions & 0 deletions tests/test_ai_storage_retention.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
"""回归:进度事件必须原地替换、按 TTL 回收并可压缩,避免 ai.sqlite3 无界膨胀。"""
import json
import sys
import threading
import time
from pathlib import Path
from types import SimpleNamespace
Expand All @@ -10,6 +11,24 @@
from wechat_decrypt_tool.local_search.service import LocalSearch


def test_periodic_maintenance_runs_until_shutdown():
from wechat_decrypt_tool.ai.lifecycle import _maintenance_loop

calls = []
stopped = threading.Event()

class Store:
def maintain(self):
calls.append(time.monotonic())
if len(calls) >= 2:
stopped.set()
return 0, 0, 0

_maintenance_loop(Store(), 'test', stopped, interval=0.01)
assert len(calls) == 2
assert calls[1] >= calls[0]


def test_local_search_update_emits_compact_deduplicated_event(tmp_path):
service = LocalSearch(tmp_path, engine=SimpleNamespace(status={}, gpu_failed=False))
job = {'id': 'job1', 'account': 'a', 'config': {'usernames': ['chat']}, 'coverage': {'0': {}},
Expand Down
Loading