diff --git a/frontend/components/LocalSearchSettings.vue b/frontend/components/LocalSearchSettings.vue index 07bc9144..f2dce5f8 100644 --- a/frontend/components/LocalSearchSettings.vue +++ b/frontend/components/LocalSearchSettings.vue @@ -378,7 +378,8 @@ function subscribe(){ if(event?.kind==='local_search_index' && event.account===current){ const incoming=event.body, existing=state.value.jobs?.find(j=>j.id===incoming?.id) if(existing && incoming.updated>=existing.updated){ - state.value.jobs=state.value.jobs.map(j=>j.id===incoming.id ? incoming : j) + // 进度事件只带高频字段;合并保留首次加载时拿到的 config/segments 等快照。 + state.value.jobs=state.value.jobs.map(j=>j.id===incoming.id ? {...j, ...incoming} : j) now.value=Date.now()/1000 if(incoming.status==='done')refresh() }else if(!existing){refresh()}else{api.diagnostic?.('response.stale',{task_id:incoming.id,component:'search'})} diff --git a/src/wechat_decrypt_tool/ai/lifecycle.py b/src/wechat_decrypt_tool/ai/lifecycle.py index cfb06d94..31c63e00 100644 --- a/src/wechat_decrypt_tool/ai/lifecycle.py +++ b/src/wechat_decrypt_tool/ai/lifecycle.py @@ -1,10 +1,21 @@ """AI 服务统一启停诊断,独立于其他后台服务的生命周期。""" import logging +import threading from importlib.metadata import PackageNotFoundError, version from .diagnostics import observed, event +def _maintain_store(store, name): + """在后台回收过期事件并压缩数据库,避免大型遗留库拖慢启动。""" + try: + deduplicated, removed, freed = store.maintain() + event('storage.maintenance.finished', component=name, events_deduplicated=deduplicated, + events_removed=removed, bytes_freed=freed) + except Exception as error: + event('storage.maintenance.failed', level=logging.WARNING, component=name, error=error) + + @observed('lifecycle.start') async def start_services(): for package in ('deepagents', 'langchain-openai', 'langchain-anthropic', 'langgraph', 'onnxruntime', 'sqlite-vec', 'tokenizers'): @@ -19,6 +30,8 @@ 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() @observed('lifecycle.stop') diff --git a/src/wechat_decrypt_tool/ai/storage.py b/src/wechat_decrypt_tool/ai/storage.py index 2998d930..7c49fca8 100644 --- a/src/wechat_decrypt_tool/ai/storage.py +++ b/src/wechat_decrypt_tool/ai/storage.py @@ -3,6 +3,7 @@ import logging import json +import os import sqlite3 import threading import time @@ -13,6 +14,28 @@ from ..app_paths import get_output_dir +# 事件默认保留窗口;超过该窗口且无需重放的事件会被回收。 +EVENT_RETENTION_SECONDS = 24 * 3600 +# 事件表空闲页超过该阈值才执行 VACUUM,避免频繁全库重写。 +COMPACT_MINIMUM_BYTES = 64 * 1024 * 1024 +# 超过该体积的库不做原地删除/VACUUM(对遗留巨型库会放大 WAL),改为重建: +# 仅保留 records 与未投递提醒,丢弃可再生的 events。 +MAINTENANCE_MAX_DATABASE_BYTES = 512 * 1024 * 1024 + +SCHEMA_SQL = """ + CREATE TABLE IF NOT EXISTS records ( + kind TEXT NOT NULL, id TEXT NOT NULL, account TEXT NOT NULL DEFAULT '', + body TEXT NOT NULL, updated REAL NOT NULL, PRIMARY KEY(kind,id)); + CREATE INDEX IF NOT EXISTS records_account ON records(kind,account,updated); + CREATE INDEX IF NOT EXISTS records_status ON records(kind,json_extract(body,'$.status'),updated); + CREATE INDEX IF NOT EXISTS records_task_usage ON records(kind,account,json_extract(body,'$.task_id')); + CREATE TABLE IF NOT EXISTS events ( + id INTEGER PRIMARY KEY AUTOINCREMENT, account TEXT NOT NULL, + kind TEXT NOT NULL, body TEXT NOT NULL, unique_key TEXT UNIQUE, + delivered INTEGER NOT NULL DEFAULT 0, created REAL NOT NULL); +""" + + class AIStore: """短事务业务存储;与工作流检查点分开,避免模型请求持有数据库锁。""" @@ -27,18 +50,7 @@ def __init__(self, root: Path | None = None): self._event_condition = threading.Condition() self._event_revisions = {} with self.connection() as db: - db.executescript(""" - CREATE TABLE IF NOT EXISTS records ( - kind TEXT NOT NULL, id TEXT NOT NULL, account TEXT NOT NULL DEFAULT '', - body TEXT NOT NULL, updated REAL NOT NULL, PRIMARY KEY(kind,id)); - CREATE INDEX IF NOT EXISTS records_account ON records(kind,account,updated); - CREATE INDEX IF NOT EXISTS records_status ON records(kind,json_extract(body,'$.status'),updated); - CREATE INDEX IF NOT EXISTS records_task_usage ON records(kind,account,json_extract(body,'$.task_id')); - CREATE TABLE IF NOT EXISTS events ( - id INTEGER PRIMARY KEY AUTOINCREMENT, account TEXT NOT NULL, - kind TEXT NOT NULL, body TEXT NOT NULL, unique_key TEXT UNIQUE, - delivered INTEGER NOT NULL DEFAULT 0, created REAL NOT NULL); - """) + db.executescript(SCHEMA_SQL) @contextmanager def connection(self): @@ -108,13 +120,30 @@ def delete(self, kind, id): with self.connection() as db: db.execute("DELETE FROM records WHERE kind=? AND id=?", (kind, id)) - def event(self, account, kind, body, unique_key=None): - inserted = False + def event(self, account, kind, body, unique_key=None, replace=False): + """写入事件供 SSE 重放。 + + 默认行为保持不变:提供 `unique_key` 时按去重语义写入(同 key 已存在则忽略), + 用于提醒等只应投递一次的事件。`replace=True` 时改为用最新快照替换旧行, + 让高频进度事件每个逻辑任务只保留一行,同时因 INSERT OR REPLACE 会删除旧行、 + 新行仍获得递增的自增 id,断线重连的 EventSource 依然能收到最新状态。 + """ with self.connection() as db: if account in self.revoked_accounts: return - cursor = db.execute("INSERT OR IGNORE INTO events(account,kind,body,unique_key,created) VALUES(?,?,?,?,?)", - (account, kind, json.dumps(body, ensure_ascii=False), unique_key, time.time())) + payload = json.dumps(body, ensure_ascii=False) + if unique_key is None: + cursor = db.execute( + "INSERT INTO events(account,kind,body,created) VALUES(?,?,?,?)", + (account, kind, payload, time.time())) + elif replace: + cursor = db.execute( + "INSERT OR REPLACE INTO events(account,kind,body,unique_key,created) VALUES(?,?,?,?,?)", + (account, kind, payload, unique_key, time.time())) + else: + cursor = db.execute( + "INSERT OR IGNORE INTO events(account,kind,body,unique_key,created) VALUES(?,?,?,?,?)", + (account, kind, payload, unique_key, time.time())) inserted = cursor.rowcount > 0 if inserted: with self._event_condition: @@ -151,6 +180,177 @@ def acknowledge(self, id): with self.connection() as db: db.execute("UPDATE events SET delivered=1 WHERE id=?", (id,)) + @observed('storage.prune_duplicates') + def prune_duplicate_events(self, batch=2000): + """一次性折叠旧版追加式进度事件:每个逻辑任务只保留最新快照。 + + 新写入的进度事件已带 unique_key、本身只保留一行;这里主要清理升级前 + 历史遗留的、同一任务多次追加的整份快照,避免巨型库只能等 TTL 慢慢过期。 + """ + total = 0 + for kind, field in (('local_search_index', '$.id'), ('local_search_download', '$.id'), + ('local_search_total', '$.job_id')): + with self.connection() as db: + db.execute("CREATE TEMP TABLE IF NOT EXISTS keep_event_ids(id INTEGER PRIMARY KEY)") + db.execute("DELETE FROM keep_event_ids") + db.execute( + f"INSERT INTO keep_event_ids SELECT max(id) FROM events " + f"WHERE kind=? AND unique_key IS NULL GROUP BY json_extract(body,'{field}')", (kind,)) + while True: + removed = db.execute( + "DELETE FROM events WHERE id IN (SELECT id FROM events " + "WHERE kind=? AND unique_key IS NULL AND id NOT IN (SELECT id FROM keep_event_ids) LIMIT ?)", + (kind, batch)).rowcount + total += removed + db.commit() + if removed < batch: + break + db.execute("DELETE FROM keep_event_ids") + for kind in ('local_search_device', 'local_search_gpu'): + with self.connection() as db: + total += db.execute( + "DELETE FROM events WHERE kind=? AND unique_key IS NULL AND id < (SELECT max(id) FROM events WHERE kind=?)", + (kind, kind)).rowcount + if total: + diagnostic_event('storage.events.deduplicated', count=total) + return total + + @observed('storage.prune_events') + def prune_events(self, max_age=EVENT_RETENTION_SECONDS, batch=2000): + """按 TTL 批量回收事件:非通知事件直接过期;已投递通知同样回收,未投递通知保留。""" + cutoff = time.time() - max_age + total = 0 + while True: + with self.connection() as db: + removed = db.execute( + "DELETE FROM events WHERE id IN (SELECT id FROM events " + "WHERE created max_database_bytes: + repaired = self.repair_oversized(max_database_bytes) + return 0, 0, repaired + deduplicated = self.prune_duplicate_events() + removed = self.prune_events(max_age) + freed = self.compact(minimum_bytes) + return deduplicated, removed, freed + + @staticmethod + def _database_bytes(db): + page_size = db.execute('PRAGMA page_size').fetchone()[0] + page_count = db.execute('PRAGMA page_count').fetchone()[0] + return page_size * page_count + @observed('storage.purge_account') def purge_account(self, account): with self.connection() as db: diff --git a/src/wechat_decrypt_tool/local_search/downloads.py b/src/wechat_decrypt_tool/local_search/downloads.py index de391361..965bd37c 100644 --- a/src/wechat_decrypt_tool/local_search/downloads.py +++ b/src/wechat_decrypt_tool/local_search/downloads.py @@ -131,7 +131,7 @@ def update(self, job, **values): job.update(values) job['updated'] = time.time() self.store.put('download', job, id=job['id']) - self.store.event('', 'local_search_download', job) + self.store.event('', 'local_search_download', job, unique_key=f'download:{job["id"]}', replace=True) if changed: diagnostic_event('download.state', level=logging.ERROR if job.get('status')=='error' else logging.INFO, model=job['id'], status=job.get('status'), phase=job.get('stage'), bytes=job.get('bytes'), attempt=job.get('attempt')) diff --git a/src/wechat_decrypt_tool/local_search/gpu.py b/src/wechat_decrypt_tool/local_search/gpu.py index 678c6f1b..2f5a3043 100644 --- a/src/wechat_decrypt_tool/local_search/gpu.py +++ b/src/wechat_decrypt_tool/local_search/gpu.py @@ -56,7 +56,7 @@ def status(self): def update(self, **values): job={**(self.store.get('gpu_component','global') or {}),**values,'updated':time.time()} self.store.put('gpu_component',job,id='global') - self.store.event('','local_search_gpu',job) + self.store.event('','local_search_gpu',job, unique_key='local_search_gpu', replace=True) @observed('gpu.start') async def start(self, source=None): diff --git a/src/wechat_decrypt_tool/local_search/service.py b/src/wechat_decrypt_tool/local_search/service.py index b4022acb..eb55a783 100644 --- a/src/wechat_decrypt_tool/local_search/service.py +++ b/src/wechat_decrypt_tool/local_search/service.py @@ -24,12 +24,15 @@ 'start': None, 'end': None, 'device': 'auto', 'device_id': 0, 'auto_update': True, 'read_batch_size': 0, 'agent_global': False} +# 进度事件里体积大且很少变化、或可从权威记录重建的字段,不随每次进度写入事件表。 +EVENT_OMITTED_FIELDS = frozenset({'config', 'coverage', 'segments', 'read_starts'}) + class LocalSearch(ProgressiveIndex, MessageTotals): def __init__(self, root=None, model_root=None, reader=None, engine=None): self.root = Path(root or get_output_dir() / 'local_search') self.store = AIStore(self.root) - self.engine = engine or LocalInference(callback=lambda status: self.store.event('', 'local_search_device', status)) + self.engine = engine or LocalInference(callback=lambda status: self.store.event('', 'local_search_device', status, unique_key='local_search_device', replace=True)) self.downloads = ModelDownloads(model_root or get_data_dir() / 'local_search_models', self.store, self.engine) from .gpu import GPUComponent self.gpu = GPUComponent(self.downloads.root.parent / 'local_search_gpu', self.store, self.engine) @@ -103,7 +106,9 @@ def update(self, job, **changes): if job['id'] in self.restarting: job['resume_on_start']=True if job['account'] in self.revoked: return self.store.put('index_job', job, id=job['id'], account=job['account']) - self.store.event(job['account'], 'local_search_index', job) + # 事件表只保留每个任务的最新进度;完整快照以 records 表为准。 + progress = {key: value for key, value in job.items() if key not in EVENT_OMITTED_FIELDS} + self.store.event(job['account'], 'local_search_index', progress, unique_key=f'index_job:{job["id"]}', replace=True) def enrichment_version(self, account): """只检查本地提取缓存,不触发媒体分析或网络访问。""" diff --git a/src/wechat_decrypt_tool/local_search/totals.py b/src/wechat_decrypt_tool/local_search/totals.py index a354f49a..1c2444c9 100644 --- a/src/wechat_decrypt_tool/local_search/totals.py +++ b/src/wechat_decrypt_tool/local_search/totals.py @@ -11,7 +11,7 @@ class MessageTotals: def save_message_total(self, job, **values): total = {'job_id': job['id'], 'updated': time.time(), **values} self.store.put('index_message_total', total, id=job['id'], account=job['account']) - self.store.event(job['account'], 'local_search_total', total) + self.store.event(job['account'], 'local_search_total', total, unique_key=f'index_total:{job["id"]}', replace=True) def message_plan(self, job): account_key = hashlib.sha256(job['account'].encode()).hexdigest() diff --git a/tests/test_ai_message_pages.py b/tests/test_ai_message_pages.py index c91e62df..4a8c0481 100644 --- a/tests/test_ai_message_pages.py +++ b/tests/test_ai_message_pages.py @@ -337,6 +337,13 @@ def encode(*args): async def run(): service = LocalSearch(tmp_path / 'state', tmp_path / 'models', engine=engine) + emitted = [] + original_event = service.store.event + def record(account, kind, body, unique_key=None, replace=False): + if kind == 'local_search_index': + emitted.append(body) + return original_event(account, kind, body, unique_key, replace) + monkeypatch.setattr(service.store, 'event', record) root = model_dir(service.downloads.root, 'bge-small-zh') root.mkdir(parents=True) Tokenizer(models.WordLevel({'[UNK]': 0}, unk_token='[UNK]')).save(str(root / 'tokenizer.json')) @@ -365,7 +372,11 @@ async def run(): 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 events) + and e['body'].get('read_count', 0) > e['body']['processed'] 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 + assert 'config' not in stored[0]['body'] and 'segments' not in stored[0]['body'] and 'coverage' not in stored[0]['body'] await service.stop() asyncio.run(run()) diff --git a/tests/test_ai_storage_retention.py b/tests/test_ai_storage_retention.py new file mode 100644 index 00000000..f6296553 --- /dev/null +++ b/tests/test_ai_storage_retention.py @@ -0,0 +1,108 @@ +"""回归:进度事件必须原地替换、按 TTL 回收并可压缩,避免 ai.sqlite3 无界膨胀。""" +import json +import sys +import time +from pathlib import Path +from types import SimpleNamespace + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) +from wechat_decrypt_tool.ai.storage import AIStore +from wechat_decrypt_tool.local_search.service import LocalSearch + + +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': {}}, + 'segments': [{}], 'read_starts': {'chat': 0}, 'processed': 0} + for processed in range(10): + service.update(job, processed=processed) + events = [event for event in service.store.events() if event['kind'] == 'local_search_index'] + assert len(events) == 1 + body = events[0]['body'] + assert body['processed'] == 9 and body['id'] == 'job1' + for heavy in ('config', 'coverage', 'segments', 'read_starts'): + assert heavy not in body + + +def test_prune_duplicate_events_collapses_legacy_snapshots(tmp_path): + store = AIStore(tmp_path) + now = time.time() + with store.connection() as db: + for processed in range(5): + db.execute("INSERT INTO events(account,kind,body,unique_key,delivered,created) VALUES('','local_search_index',?,NULL,0,?)", + (json.dumps({'id': 'job-a', 'processed': processed}), now)) + for processed in (1, 2): + db.execute("INSERT INTO events(account,kind,body,unique_key,delivered,created) VALUES('','local_search_index',?,NULL,0,?)", + (json.dumps({'id': 'job-b', 'processed': processed}), now)) + assert store.prune_duplicate_events(batch=2) == 5 + kept = sorted(row['body']['processed'] for row in store.events()) + assert kept == [2, 4] + + +def test_unique_key_event_keeps_only_latest_snapshot(tmp_path): + store = AIStore(tmp_path) + for processed in range(50): + store.event('', 'local_search_index', {'id': 'job', 'processed': processed}, + unique_key='index_job:job', replace=True) + rows = store.events() + assert len(rows) == 1 + assert rows[0]['body']['processed'] == 49 + # 替换后仍产生新的自增 id,保证 SSE 客户端能收到更新。 + assert rows[0]['id'] > 0 + + +def test_unique_key_without_replace_stays_idempotent(tmp_path): + store = AIStore(tmp_path) + store.event('', 'notification', {'n': 1}, unique_key='summary:1') + store.event('', 'notification', {'n': 2}, unique_key='summary:1') + rows = store.events() + assert len(rows) == 1 and rows[0]['body']['n'] == 1 + + +def test_prune_events_keeps_undelivered_notifications(tmp_path): + store = AIStore(tmp_path) + old = time.time() - 10 * 24 * 3600 + with store.connection() as db: + db.execute("INSERT INTO events(account,kind,body,unique_key,delivered,created) VALUES('','local_search_index','{}','old-progress',0,?)", (old,)) + db.execute("INSERT INTO events(account,kind,body,unique_key,delivered,created) VALUES('','notification','{}',NULL,0,?)", (old,)) + db.execute("INSERT INTO events(account,kind,body,unique_key,delivered,created) VALUES('','notification','{}',NULL,1,?)", (old,)) + assert store.prune_events(max_age=24 * 3600) == 2 + remaining = store.events() + assert len(remaining) == 1 + assert remaining[0]['kind'] == 'notification' and remaining[0]['delivered'] == 0 + + +def test_maintain_repairs_oversized_database_preserving_records(tmp_path): + store = AIStore(tmp_path) + store.put('config', {'enabled': True}, id='acct', account='acct') + now = time.time() + with store.connection() as db: + for index in range(300): + db.execute("INSERT INTO events(account,kind,body,unique_key,delivered,created) VALUES('','local_search_index',?,NULL,0,?)", + (json.dumps({'id': 'job', 'index': index, 'pad': 'x' * 2000}), now)) + db.execute("INSERT INTO events(account,kind,body,unique_key,delivered,created) VALUES('','notification','{\"n\":1}','k1',0,?)", (now,)) + db.execute("INSERT INTO events(account,kind,body,unique_key,delivered,created) VALUES('','notification','{\"n\":2}','k2',1,?)", (now,)) + before = store._database_bytes(db) + # 用极小阈值模拟遗留巨型库:应重建而不是长时间原地删除。 + deduplicated, pruned, repaired = store.maintain(max_database_bytes=1) + assert repaired > 0 + with store.connection() as db: + assert store._database_bytes(db) < before + # records 保留;未投递提醒保留,可再生的进度事件与已投递提醒丢弃。 + assert store.get('config', 'acct')['enabled'] is True + events = store.events() + assert [event['kind'] for event in events] == ['notification'] + assert events[0]['body']['n'] == 1 + + +def test_compact_reclaims_free_pages_after_prune(tmp_path): + store = AIStore(tmp_path) + for index in range(2000): + store.event('', 'task', {'payload': 'x' * 2000, 'index': index}) + with store.connection() as db: + db.execute("DELETE FROM events") + before = store.path.stat().st_size + assert before > 0 + freed = store.compact(minimum_bytes=1) + assert freed > 0 + assert store.path.stat().st_size < before