feat(memory): implement OpenRouter + pgvector memory pipeline

2026-02-25 16:21:50 +08:00
parent 049ed35557
commit 40330ac90d
15 changed files with 1744 additions and 0 deletions
@@ -0,0 +1,108 @@
+#!/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()