109 lines
4.1 KiB
Python
Executable File
109 lines
4.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
from pathlib import Path
|
|
|
|
from index_common import (
|
|
VAULT_ROOT,
|
|
embed_text,
|
|
get_conn,
|
|
index_lock,
|
|
is_excluded,
|
|
normalize_rel,
|
|
sha256_text,
|
|
vector_literal,
|
|
)
|
|
|
|
PRIMARY_DIRS = ['01_Projects', '02_Areas']
|
|
MIN_TEXT_LEN = 50
|
|
|
|
|
|
def upsert_primary(cur, rel: str, text: str, embedding: list[float]) -> None:
|
|
cur.execute(
|
|
'''
|
|
INSERT INTO memory_primary (id, source, content, content_hash, embedding)
|
|
VALUES (%s, %s, %s, %s, %s::vector)
|
|
ON CONFLICT (id) DO UPDATE SET
|
|
source = EXCLUDED.source,
|
|
content = EXCLUDED.content,
|
|
content_hash = EXCLUDED.content_hash,
|
|
embedding = EXCLUDED.embedding,
|
|
updated_at = now()
|
|
''',
|
|
(rel, rel, text, sha256_text(text), vector_literal(embedding)),
|
|
)
|
|
|
|
|
|
def run() -> None:
|
|
with index_lock():
|
|
with get_conn() as conn:
|
|
with conn.cursor() as cur:
|
|
valid_ids: set[str] = set()
|
|
secure_ids: set[str] = set()
|
|
failed_ids: list[str] = []
|
|
|
|
print(f'[INFO] 开始全量扫描: {VAULT_ROOT}')
|
|
for dir_name in PRIMARY_DIRS:
|
|
target = VAULT_ROOT / dir_name
|
|
if not target.exists():
|
|
continue
|
|
|
|
for md_file in target.rglob('*.md'):
|
|
rel = normalize_rel(md_file)
|
|
if is_excluded(md_file):
|
|
secure_ids.add(rel)
|
|
cur.execute(
|
|
'''
|
|
INSERT INTO memory_secure_audit (id, source, risk)
|
|
VALUES (%s, %s, %s)
|
|
ON CONFLICT (id) DO UPDATE SET
|
|
risk = EXCLUDED.risk,
|
|
updated_at = now()
|
|
''',
|
|
(rel, rel, 'excluded_or_sensitive'),
|
|
)
|
|
cur.execute('DELETE FROM memory_primary WHERE id=%s', (rel,))
|
|
print(f'[QUARANTINED] {rel}')
|
|
continue
|
|
|
|
text = md_file.read_text(encoding='utf-8', errors='ignore')
|
|
if len(text.strip()) < MIN_TEXT_LEN:
|
|
cur.execute('DELETE FROM memory_primary WHERE id=%s', (rel,))
|
|
cur.execute('DELETE FROM memory_secure_audit WHERE id=%s', (rel,))
|
|
print(f'[PRUNED] {rel}')
|
|
continue
|
|
|
|
try:
|
|
embedding = embed_text(text)
|
|
except Exception as error:
|
|
failed_ids.append(rel)
|
|
print(f'[EMBED-ERROR] {rel} :: {error}')
|
|
continue
|
|
upsert_primary(cur, rel, text, embedding)
|
|
cur.execute('DELETE FROM memory_secure_audit WHERE id=%s', (rel,))
|
|
valid_ids.add(rel)
|
|
print(f'[UPSERTED] {rel}')
|
|
|
|
cur.execute('SELECT id FROM memory_primary')
|
|
db_primary_ids = {row[0] for row in cur.fetchall()}
|
|
stale_primary_ids = sorted(db_primary_ids - valid_ids)
|
|
for stale_id in stale_primary_ids:
|
|
cur.execute('DELETE FROM memory_primary WHERE id=%s', (stale_id,))
|
|
|
|
cur.execute('SELECT id FROM memory_secure_audit')
|
|
db_secure_ids = {row[0] for row in cur.fetchall()}
|
|
stale_secure_ids = sorted(db_secure_ids - secure_ids)
|
|
for stale_id in stale_secure_ids:
|
|
cur.execute('DELETE FROM memory_secure_audit WHERE id=%s', (stale_id,))
|
|
|
|
conn.commit()
|
|
|
|
print(
|
|
f'[DONE] 全量索引完成 primary={len(valid_ids)} secure={len(secure_ids)} '
|
|
f'clean_primary={len(stale_primary_ids)} clean_secure={len(stale_secure_ids)} '
|
|
f'failed={len(failed_ids)}'
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
run()
|