import json import re from datetime import datetime, timedelta from typing import Any, Dict, Iterable, List, Optional, Sequence from config import POSTGRES_CONNECTION_STRING, WIKI_CONFIG async def _connect_pool(): from psycopg_pool import AsyncConnectionPool return AsyncConnectionPool( POSTGRES_CONNECTION_STRING, kwargs={"autocommit": True}, ) async def init_wiki_tables() -> None: """Create Wiki pointer tables in the wx-agent PostgreSQL database.""" try: pool = await _connect_pool() except ImportError: print("[wiki_engine] psycopg_pool is not installed; skip wiki table init") return async with pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute(""" CREATE TABLE IF NOT EXISTS wiki_sync_jobs ( id TEXT PRIMARY KEY, target_type TEXT NOT NULL DEFAULT '', target_id TEXT NOT NULL DEFAULT '', kb_id TEXT NOT NULL DEFAULT '', status TEXT NOT NULL DEFAULT 'pending', total_files INTEGER NOT NULL DEFAULT 0, processed_files INTEGER NOT NULL DEFAULT 0, total_slices INTEGER NOT NULL DEFAULT 0, skipped_files INTEGER NOT NULL DEFAULT 0, page_count INTEGER NOT NULL DEFAULT 0, stage TEXT NOT NULL DEFAULT '', current_file_id TEXT NOT NULL DEFAULT '', current_filename TEXT NOT NULL DEFAULT '', error TEXT NOT NULL DEFAULT '', created_at TIMESTAMP NOT NULL DEFAULT NOW(), updated_at TIMESTAMP NOT NULL DEFAULT NOW() ) """) await cur.execute("ALTER TABLE wiki_sync_jobs ADD COLUMN IF NOT EXISTS skipped_files INTEGER NOT NULL DEFAULT 0") await cur.execute("ALTER TABLE wiki_sync_jobs ADD COLUMN IF NOT EXISTS page_count INTEGER NOT NULL DEFAULT 0") await cur.execute("ALTER TABLE wiki_sync_jobs ADD COLUMN IF NOT EXISTS stage TEXT NOT NULL DEFAULT ''") await cur.execute("ALTER TABLE wiki_sync_jobs ADD COLUMN IF NOT EXISTS current_file_id TEXT NOT NULL DEFAULT ''") await cur.execute("ALTER TABLE wiki_sync_jobs ADD COLUMN IF NOT EXISTS current_filename TEXT NOT NULL DEFAULT ''") await cur.execute(""" CREATE TABLE IF NOT EXISTS wiki_file_state ( id SERIAL PRIMARY KEY, kb_id TEXT NOT NULL DEFAULT '', file_id TEXT NOT NULL UNIQUE, filename TEXT NOT NULL DEFAULT '', file_hash TEXT NOT NULL DEFAULT '', status TEXT NOT NULL DEFAULT '', slice_count INTEGER NOT NULL DEFAULT 0, source_updated_at TEXT NOT NULL DEFAULT '', metadata JSONB NOT NULL DEFAULT '{}', last_synced_at TIMESTAMP NOT NULL DEFAULT NOW() ) """) await cur.execute(""" CREATE TABLE IF NOT EXISTS wiki_slice_refs ( id SERIAL PRIMARY KEY, kb_id TEXT NOT NULL DEFAULT '', file_id TEXT NOT NULL DEFAULT '', slice_id TEXT NOT NULL UNIQUE, ordinal INTEGER NOT NULL DEFAULT 0, filename TEXT NOT NULL DEFAULT '', title_path JSONB NOT NULL DEFAULT '[]', keywords JSONB NOT NULL DEFAULT '[]', brief TEXT NOT NULL DEFAULT '', content TEXT NOT NULL DEFAULT '', chunk_type TEXT NOT NULL DEFAULT 'text', content_hash TEXT NOT NULL DEFAULT '', updated_at TIMESTAMP NOT NULL DEFAULT NOW() ) """) await cur.execute("ALTER TABLE wiki_slice_refs ADD COLUMN IF NOT EXISTS content TEXT NOT NULL DEFAULT ''") await cur.execute(""" CREATE TABLE IF NOT EXISTS wiki_pages ( id SERIAL PRIMARY KEY, kb_id TEXT NOT NULL DEFAULT '', slug TEXT NOT NULL DEFAULT '', title TEXT NOT NULL DEFAULT '', page_type TEXT NOT NULL DEFAULT 'topic', summary TEXT NOT NULL DEFAULT '', content TEXT NOT NULL DEFAULT '', keywords JSONB NOT NULL DEFAULT '[]', aliases JSONB NOT NULL DEFAULT '[]', related_file_ids JSONB NOT NULL DEFAULT '[]', source_refs JSONB NOT NULL DEFAULT '[]', chunk_refs JSONB NOT NULL DEFAULT '[]', source_kind TEXT NOT NULL DEFAULT 'rule', score_hint REAL NOT NULL DEFAULT 1.0, created_at TIMESTAMP NOT NULL DEFAULT NOW(), updated_at TIMESTAMP NOT NULL DEFAULT NOW(), UNIQUE(kb_id, title, page_type) ) """) await cur.execute("ALTER TABLE wiki_pages ADD COLUMN IF NOT EXISTS slug TEXT NOT NULL DEFAULT ''") await cur.execute("ALTER TABLE wiki_pages ADD COLUMN IF NOT EXISTS content TEXT NOT NULL DEFAULT ''") await cur.execute("ALTER TABLE wiki_pages ADD COLUMN IF NOT EXISTS aliases JSONB NOT NULL DEFAULT '[]'") await cur.execute("ALTER TABLE wiki_pages ADD COLUMN IF NOT EXISTS source_refs JSONB NOT NULL DEFAULT '[]'") await cur.execute("ALTER TABLE wiki_pages ADD COLUMN IF NOT EXISTS chunk_refs JSONB NOT NULL DEFAULT '[]'") await cur.execute(""" CREATE TABLE IF NOT EXISTS wiki_page_slices ( id SERIAL PRIMARY KEY, wiki_page_id INTEGER NOT NULL REFERENCES wiki_pages(id) ON DELETE CASCADE, file_id TEXT NOT NULL DEFAULT '', slice_id TEXT NOT NULL DEFAULT '', weight REAL NOT NULL DEFAULT 1.0, reason TEXT NOT NULL DEFAULT '', UNIQUE(wiki_page_id, slice_id) ) """) await cur.execute(""" CREATE TABLE IF NOT EXISTS wiki_query_logs ( id SERIAL PRIMARY KEY, query TEXT NOT NULL DEFAULT '', kb_id TEXT NOT NULL DEFAULT '', result_count INTEGER NOT NULL DEFAULT 0, created_at TIMESTAMP NOT NULL DEFAULT NOW() ) """) await cur.execute(""" CREATE TABLE IF NOT EXISTS wiki_pending_ops ( id BIGSERIAL PRIMARY KEY, job_id TEXT NOT NULL DEFAULT '', kb_id TEXT NOT NULL DEFAULT '', file_id TEXT NOT NULL DEFAULT '', op TEXT NOT NULL DEFAULT 'sync_file', dedup_key TEXT NOT NULL DEFAULT '', payload JSONB NOT NULL DEFAULT '{}', status TEXT NOT NULL DEFAULT 'pending', fail_count INTEGER NOT NULL DEFAULT 0, error TEXT NOT NULL DEFAULT '', claimed_at TIMESTAMP, enqueued_at TIMESTAMP NOT NULL DEFAULT NOW(), updated_at TIMESTAMP NOT NULL DEFAULT NOW() ) """) await cur.execute(""" CREATE TABLE IF NOT EXISTS wiki_dead_letters ( id BIGSERIAL PRIMARY KEY, job_id TEXT NOT NULL DEFAULT '', kb_id TEXT NOT NULL DEFAULT '', file_id TEXT NOT NULL DEFAULT '', op TEXT NOT NULL DEFAULT '', payload JSONB NOT NULL DEFAULT '{}', last_error TEXT NOT NULL DEFAULT '', fail_count INTEGER NOT NULL DEFAULT 0, failed_at TIMESTAMP NOT NULL DEFAULT NOW() ) """) await cur.execute(""" CREATE TABLE IF NOT EXISTS wiki_llm_cache ( cache_key TEXT PRIMARY KEY, response TEXT NOT NULL DEFAULT '', hit_count INTEGER NOT NULL DEFAULT 0, created_at TIMESTAMP NOT NULL DEFAULT NOW(), updated_at TIMESTAMP NOT NULL DEFAULT NOW() ) """) await cur.execute(""" CREATE TABLE IF NOT EXISTS wiki_page_locks ( kb_id TEXT NOT NULL DEFAULT '', slug TEXT NOT NULL DEFAULT '', owner TEXT NOT NULL DEFAULT '', locked_at TIMESTAMP NOT NULL DEFAULT NOW(), expires_at TIMESTAMP NOT NULL DEFAULT NOW(), PRIMARY KEY (kb_id, slug) ) """) await cur.execute(""" CREATE TABLE IF NOT EXISTS wiki_finalize_ops ( id BIGSERIAL PRIMARY KEY, job_id TEXT NOT NULL DEFAULT '', kb_id TEXT NOT NULL DEFAULT '', status TEXT NOT NULL DEFAULT 'pending', payload JSONB NOT NULL DEFAULT '{}', claimed_at TIMESTAMP, run_after TIMESTAMP NOT NULL DEFAULT NOW(), enqueued_at TIMESTAMP NOT NULL DEFAULT NOW(), updated_at TIMESTAMP NOT NULL DEFAULT NOW() ) """) await cur.execute(""" CREATE TABLE IF NOT EXISTS wiki_admin_plans ( id TEXT PRIMARY KEY, chat_id TEXT NOT NULL DEFAULT '', status TEXT NOT NULL DEFAULT 'pending', intent TEXT NOT NULL DEFAULT '', scope TEXT NOT NULL DEFAULT '', kb_id TEXT NOT NULL DEFAULT '', kb_name TEXT NOT NULL DEFAULT '', summary JSONB NOT NULL DEFAULT '{}', actions JSONB NOT NULL DEFAULT '[]', executed JSONB NOT NULL DEFAULT '[]', metadata JSONB NOT NULL DEFAULT '{}', created_at TIMESTAMP NOT NULL DEFAULT NOW(), updated_at TIMESTAMP NOT NULL DEFAULT NOW(), expires_at TIMESTAMP NOT NULL DEFAULT NOW() ) """) await cur.execute("CREATE INDEX IF NOT EXISTS wiki_file_state_kb_idx ON wiki_file_state(kb_id)") await cur.execute("CREATE INDEX IF NOT EXISTS wiki_slice_refs_kb_file_idx ON wiki_slice_refs(kb_id, file_id)") await cur.execute("CREATE INDEX IF NOT EXISTS wiki_slice_refs_file_idx ON wiki_slice_refs(file_id)") await cur.execute("CREATE INDEX IF NOT EXISTS wiki_pages_kb_idx ON wiki_pages(kb_id)") await cur.execute("CREATE INDEX IF NOT EXISTS wiki_pages_title_idx ON wiki_pages(title)") await cur.execute("CREATE INDEX IF NOT EXISTS wiki_pages_slug_idx ON wiki_pages(slug)") await cur.execute("CREATE INDEX IF NOT EXISTS wiki_page_slices_slice_idx ON wiki_page_slices(slice_id)") await cur.execute("CREATE INDEX IF NOT EXISTS wiki_pending_ops_status_idx ON wiki_pending_ops(status, claimed_at, id)") await cur.execute("CREATE INDEX IF NOT EXISTS wiki_pending_ops_kb_idx ON wiki_pending_ops(kb_id, status, id)") await cur.execute("CREATE UNIQUE INDEX IF NOT EXISTS wiki_pending_ops_job_file_op_idx ON wiki_pending_ops(job_id, file_id, op)") await cur.execute("CREATE INDEX IF NOT EXISTS wiki_dead_letters_kb_idx ON wiki_dead_letters(kb_id, failed_at DESC)") await cur.execute("CREATE INDEX IF NOT EXISTS wiki_llm_cache_updated_idx ON wiki_llm_cache(updated_at)") await cur.execute("CREATE INDEX IF NOT EXISTS wiki_page_locks_expire_idx ON wiki_page_locks(expires_at)") await cur.execute("CREATE INDEX IF NOT EXISTS wiki_finalize_ops_status_idx ON wiki_finalize_ops(status, run_after, id)") await cur.execute("CREATE INDEX IF NOT EXISTS wiki_finalize_ops_kb_idx ON wiki_finalize_ops(kb_id, status)") await cur.execute("CREATE INDEX IF NOT EXISTS wiki_admin_plans_chat_status_idx ON wiki_admin_plans(chat_id, status, created_at DESC)") print("[wiki_engine] wiki tables initialized") async def create_job(job_id: str, target_type: str, target_id: str = "", kb_id: str = "") -> None: async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ INSERT INTO wiki_sync_jobs (id, target_type, target_id, kb_id, status, created_at, updated_at) VALUES (%s, %s, %s, %s, 'running', %s, %s) ON CONFLICT (id) DO UPDATE SET status = EXCLUDED.status, updated_at = EXCLUDED.updated_at """, (job_id, target_type, target_id or "", kb_id or "", datetime.now(), datetime.now()), ) async def update_job(job_id: str, **fields: Any) -> None: if not fields: return allowed = { "status", "total_files", "processed_files", "total_slices", "error", "skipped_files", "page_count", "stage", "current_file_id", "current_filename", } parts = [] params: List[Any] = [] for key, value in fields.items(): if key in allowed: parts.append(f"{key} = %s") params.append(value) if not parts: return parts.append("updated_at = %s") params.append(datetime.now()) params.append(job_id) async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute(f"UPDATE wiki_sync_jobs SET {', '.join(parts)} WHERE id = %s", params) async def get_job(job_id: str) -> Optional[Dict[str, Any]]: async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ SELECT id, target_type, target_id, kb_id, status, total_files, processed_files, total_slices, skipped_files, page_count, stage, current_file_id, current_filename, error, created_at, updated_at FROM wiki_sync_jobs WHERE id = %s """, (job_id,), ) row = await cur.fetchone() if not row: return None keys = [ "id", "target_type", "target_id", "kb_id", "status", "total_files", "processed_files", "total_slices", "skipped_files", "page_count", "stage", "current_file_id", "current_filename", "error", "created_at", "updated_at", ] return _row_to_dict(keys, row) async def set_job_totals(job_id: str, total_files: int) -> None: await update_job( job_id, total_files=max(int(total_files or 0), 0), processed_files=0, total_slices=0, skipped_files=0, page_count=0, stage="queued", current_file_id="", current_filename="", ) async def increment_job_progress( job_id: str, processed_delta: int = 1, slice_delta: int = 0, skipped_delta: int = 0, page_delta: int = 0, error: str = "", stage: str = "", current_file_id: str = "", current_filename: str = "", ) -> None: processed_delta = max(int(processed_delta or 0), 0) slice_delta = max(int(slice_delta or 0), 0) skipped_delta = max(int(skipped_delta or 0), 0) page_delta = max(int(page_delta or 0), 0) error = str(error or "").strip() async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ UPDATE wiki_sync_jobs SET processed_files = processed_files + %s, total_slices = total_slices + %s, skipped_files = skipped_files + %s, page_count = page_count + %s, stage = CASE WHEN %s <> '' THEN %s ELSE stage END, current_file_id = CASE WHEN %s <> '' THEN %s ELSE current_file_id END, current_filename = CASE WHEN %s <> '' THEN %s ELSE current_filename END, error = CASE WHEN %s <> '' THEN LEFT( CASE WHEN error <> '' THEN error || E'\n' || %s ELSE %s END, 8000 ) ELSE error END, status = CASE WHEN processed_files + %s >= total_files THEN CASE WHEN error <> '' OR %s <> '' THEN 'failed' ELSE 'success' END ELSE status END, updated_at = %s WHERE id = %s """, ( processed_delta, slice_delta, skipped_delta, page_delta, str(stage or ""), str(stage or ""), str(current_file_id or ""), str(current_file_id or ""), str(current_filename or ""), str(current_filename or ""), error, error, error, processed_delta, error, datetime.now(), job_id, ), ) async def enqueue_wiki_file_op( job_id: str, kb_id: str, file_id: str, file_info: Optional[Dict[str, Any]] = None, use_llm: bool = False, force: bool = False, ) -> None: payload = { "job_id": str(job_id or ""), "kb_id": str(kb_id or ""), "file_id": str(file_id or ""), "file_info": file_info or {}, "use_llm": bool(use_llm), "force": bool(force), } async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ INSERT INTO wiki_pending_ops ( job_id, kb_id, file_id, op, dedup_key, payload, status, fail_count, error, claimed_at, enqueued_at, updated_at ) VALUES (%s, %s, %s, 'sync_file', %s, %s::jsonb, 'pending', 0, '', NULL, %s, %s) ON CONFLICT (job_id, file_id, op) DO UPDATE SET kb_id = EXCLUDED.kb_id, dedup_key = EXCLUDED.dedup_key, payload = EXCLUDED.payload, status = 'pending', error = '', claimed_at = NULL, updated_at = EXCLUDED.updated_at """, ( str(job_id or ""), str(kb_id or ""), str(file_id or ""), str(file_id or ""), json.dumps(payload, ensure_ascii=False), datetime.now(), datetime.now(), ), ) async def claim_wiki_file_ops( limit: int = 1, stale_seconds: int = 1800, per_kb_concurrency: int = 1, ) -> List[Dict[str, Any]]: limit = max(1, min(int(limit or 1), 50)) stale_seconds = max(int(stale_seconds or 1800), 60) per_kb_concurrency = max(int(per_kb_concurrency or 1), 1) stale_before = datetime.now() - timedelta(seconds=stale_seconds) async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.transaction(): async with conn.cursor() as cur: await cur.execute( """ SELECT id FROM wiki_pending_ops q WHERE q.op = 'sync_file' AND ( q.status = 'pending' OR (q.status = 'running' AND (q.claimed_at IS NULL OR q.claimed_at < %s)) ) AND ( SELECT COUNT(*) FROM wiki_pending_ops active WHERE active.kb_id = q.kb_id AND active.status = 'running' AND active.claimed_at IS NOT NULL AND active.claimed_at >= %s AND active.id <> q.id ) < %s AND NOT EXISTS ( SELECT 1 FROM wiki_pending_ops active_file WHERE active_file.file_id = q.file_id AND active_file.status = 'running' AND active_file.claimed_at IS NOT NULL AND active_file.claimed_at >= %s AND active_file.id <> q.id ) ORDER BY q.id ASC LIMIT %s FOR UPDATE SKIP LOCKED """, (stale_before, stale_before, per_kb_concurrency, stale_before, limit), ) ids = [int(row[0]) for row in await cur.fetchall()] if not ids: return [] await cur.execute( """ UPDATE wiki_pending_ops SET status = 'running', claimed_at = %s, updated_at = %s WHERE id = ANY(%s) RETURNING id, job_id, kb_id, file_id, op, payload, fail_count """, (datetime.now(), datetime.now(), ids), ) rows = await cur.fetchall() keys = ["id", "job_id", "kb_id", "file_id", "op", "payload", "fail_count"] return [_row_to_dict(keys, row) for row in rows] async def complete_wiki_op(op_id: int) -> None: async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute("DELETE FROM wiki_pending_ops WHERE id = %s", (int(op_id),)) async def fail_wiki_op(op: Dict[str, Any], error: str, max_retries: int) -> bool: op_id = int(op.get("id") or 0) max_retries = max(int(max_retries or 1), 1) error = str(error or "")[:4000] async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ UPDATE wiki_pending_ops SET fail_count = fail_count + 1, error = %s, status = CASE WHEN fail_count + 1 >= %s THEN 'dead' ELSE 'pending' END, claimed_at = NULL, updated_at = %s WHERE id = %s RETURNING fail_count, status, payload """, (error, max_retries, datetime.now(), op_id), ) row = await cur.fetchone() if not row: return True fail_count, status, payload = row if status == "dead": await cur.execute( """ INSERT INTO wiki_dead_letters ( job_id, kb_id, file_id, op, payload, last_error, fail_count, failed_at ) VALUES (%s, %s, %s, %s, %s::jsonb, %s, %s, %s) """, ( str(op.get("job_id") or ""), str(op.get("kb_id") or ""), str(op.get("file_id") or ""), str(op.get("op") or ""), json.dumps(payload or {}, ensure_ascii=False), error, int(fail_count or 0), datetime.now(), ), ) await cur.execute("DELETE FROM wiki_pending_ops WHERE id = %s", (op_id,)) return True return False async def pending_wiki_op_count(job_id: Optional[str] = None, kb_id: Optional[str] = None) -> int: where = ["status IN ('pending', 'running')"] params: List[Any] = [] if job_id: where.append("job_id = %s") params.append(str(job_id)) if kb_id: where.append("kb_id = %s") params.append(str(kb_id)) async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute(f"SELECT COUNT(*) FROM wiki_pending_ops WHERE {' AND '.join(where)}", params) row = await cur.fetchone() return int(row[0] if row else 0) async def get_wiki_queue_stats() -> Dict[str, Any]: async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ SELECT status, COUNT(*) FROM wiki_pending_ops GROUP BY status """ ) rows = await cur.fetchall() await cur.execute("SELECT COUNT(*) FROM wiki_dead_letters") dead = (await cur.fetchone())[0] stats = {str(status): int(count) for status, count in rows} stats["dead_letters"] = int(dead or 0) return stats async def get_llm_cache(cache_key: str, ttl_seconds: int) -> Optional[str]: if not cache_key: return None ttl_seconds = int(ttl_seconds or 0) where = "cache_key = %s" params: List[Any] = [cache_key] if ttl_seconds > 0: where += " AND updated_at >= NOW() - (%s * INTERVAL '1 second')" params.append(ttl_seconds) async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute(f"SELECT response FROM wiki_llm_cache WHERE {where}", params) row = await cur.fetchone() if not row: return None await cur.execute( "UPDATE wiki_llm_cache SET hit_count = hit_count + 1, updated_at = %s WHERE cache_key = %s", (datetime.now(), cache_key), ) return str(row[0] or "") async def set_llm_cache(cache_key: str, response: str) -> None: if not cache_key: return async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ INSERT INTO wiki_llm_cache (cache_key, response, hit_count, created_at, updated_at) VALUES (%s, %s, 0, %s, %s) ON CONFLICT (cache_key) DO UPDATE SET response = EXCLUDED.response, updated_at = EXCLUDED.updated_at """, (cache_key, response or "", datetime.now(), datetime.now()), ) async def create_admin_plan( plan_id: str, chat_id: str, intent: str, scope: str, kb_id: str = "", kb_name: str = "", summary: Optional[Dict[str, Any]] = None, actions: Optional[List[Dict[str, Any]]] = None, metadata: Optional[Dict[str, Any]] = None, ttl_hours: int = 24, ) -> Dict[str, Any]: now = datetime.now() expires_at = now + timedelta(hours=max(int(ttl_hours or 24), 1)) async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ INSERT INTO wiki_admin_plans ( id, chat_id, status, intent, scope, kb_id, kb_name, summary, actions, executed, metadata, created_at, updated_at, expires_at ) VALUES (%s, %s, 'pending', %s, %s, %s, %s, %s::jsonb, %s::jsonb, '[]'::jsonb, %s::jsonb, %s, %s, %s) ON CONFLICT (id) DO UPDATE SET chat_id = EXCLUDED.chat_id, status = EXCLUDED.status, intent = EXCLUDED.intent, scope = EXCLUDED.scope, kb_id = EXCLUDED.kb_id, kb_name = EXCLUDED.kb_name, summary = EXCLUDED.summary, actions = EXCLUDED.actions, executed = EXCLUDED.executed, metadata = EXCLUDED.metadata, updated_at = EXCLUDED.updated_at, expires_at = EXCLUDED.expires_at RETURNING id, chat_id, status, intent, scope, kb_id, kb_name, summary, actions, executed, metadata, created_at, updated_at, expires_at """, ( plan_id, chat_id or "", intent or "", scope or "", kb_id or "", kb_name or "", json.dumps(summary or {}, ensure_ascii=False), json.dumps(actions or [], ensure_ascii=False), json.dumps(metadata or {}, ensure_ascii=False), now, now, expires_at, ), ) row = await cur.fetchone() return _admin_plan_from_row(row) async def get_latest_pending_admin_plan(chat_id: str) -> Optional[Dict[str, Any]]: if not chat_id: return None async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ SELECT id, chat_id, status, intent, scope, kb_id, kb_name, summary, actions, executed, metadata, created_at, updated_at, expires_at FROM wiki_admin_plans WHERE chat_id = %s AND status = 'pending' AND expires_at >= NOW() ORDER BY created_at DESC LIMIT 1 """, (chat_id,), ) row = await cur.fetchone() return _admin_plan_from_row(row) if row else None async def update_admin_plan_status( plan_id: str, status: str, executed: Optional[List[Dict[str, Any]]] = None, ) -> Optional[Dict[str, Any]]: async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: if executed is None: await cur.execute( """ UPDATE wiki_admin_plans SET status = %s, updated_at = %s WHERE id = %s RETURNING id, chat_id, status, intent, scope, kb_id, kb_name, summary, actions, executed, metadata, created_at, updated_at, expires_at """, (status, datetime.now(), plan_id), ) else: await cur.execute( """ UPDATE wiki_admin_plans SET status = %s, executed = %s::jsonb, updated_at = %s WHERE id = %s RETURNING id, chat_id, status, intent, scope, kb_id, kb_name, summary, actions, executed, metadata, created_at, updated_at, expires_at """, (status, json.dumps(executed, ensure_ascii=False), datetime.now(), plan_id), ) row = await cur.fetchone() return _admin_plan_from_row(row) if row else None async def upsert_file_state(file_info: Dict[str, Any], slice_count: int) -> None: metadata = dict(file_info or {}) nested_metadata = file_info.get("metadata") if isinstance(file_info.get("metadata"), dict) else {} metadata.update(nested_metadata) async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ INSERT INTO wiki_file_state ( kb_id, file_id, filename, file_hash, status, slice_count, source_updated_at, metadata, last_synced_at ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s::jsonb, %s) ON CONFLICT (file_id) DO UPDATE SET kb_id = EXCLUDED.kb_id, filename = EXCLUDED.filename, file_hash = EXCLUDED.file_hash, status = EXCLUDED.status, slice_count = EXCLUDED.slice_count, source_updated_at = EXCLUDED.source_updated_at, metadata = EXCLUDED.metadata, last_synced_at = EXCLUDED.last_synced_at """, ( str(file_info.get("kb_id") or ""), str(file_info.get("file_id") or file_info.get("id") or ""), str(file_info.get("filename") or file_info.get("name") or ""), str(file_info.get("hash") or ""), str(file_info.get("status") or ""), slice_count, str(file_info.get("updated_at") or ""), json.dumps(metadata, ensure_ascii=False), datetime.now(), ), ) async def get_file_state(file_id: str) -> Optional[Dict[str, Any]]: async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ SELECT kb_id, file_id, filename, file_hash, status, slice_count, source_updated_at, metadata, last_synced_at FROM wiki_file_state WHERE file_id = %s """, (str(file_id or ""),), ) row = await cur.fetchone() if not row: return None keys = [ "kb_id", "file_id", "filename", "file_hash", "status", "slice_count", "source_updated_at", "metadata", "last_synced_at", ] return _row_to_dict(keys, row) async def count_file_pages(file_id: str) -> int: async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( "SELECT COUNT(DISTINCT wiki_page_id) FROM wiki_page_slices WHERE file_id = %s", (str(file_id or ""),), ) row = await cur.fetchone() return int(row[0] or 0) if row else 0 async def replace_file_slice_refs(file_id: str, refs: Sequence[Dict[str, Any]]) -> None: async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute("DELETE FROM wiki_page_slices WHERE file_id = %s", (str(file_id),)) await cur.execute("DELETE FROM wiki_slice_refs WHERE file_id = %s", (str(file_id),)) await _delete_optional_table_rows(cur, "wiki_slice_embeddings", "file_id = %s", [str(file_id)]) for ref in refs: await cur.execute( """ INSERT INTO wiki_slice_refs ( kb_id, file_id, slice_id, ordinal, filename, title_path, keywords, brief, content, chunk_type, content_hash, updated_at ) VALUES (%s, %s, %s, %s, %s, %s::jsonb, %s::jsonb, %s, %s, %s, %s, %s) ON CONFLICT (slice_id) DO UPDATE SET kb_id = EXCLUDED.kb_id, file_id = EXCLUDED.file_id, ordinal = EXCLUDED.ordinal, filename = EXCLUDED.filename, title_path = EXCLUDED.title_path, keywords = EXCLUDED.keywords, brief = EXCLUDED.brief, content = EXCLUDED.content, chunk_type = EXCLUDED.chunk_type, content_hash = EXCLUDED.content_hash, updated_at = EXCLUDED.updated_at """, ( str(ref.get("kb_id") or ""), str(ref.get("file_id") or ""), str(ref.get("slice_id") or ""), int(ref.get("ordinal") or 0), str(ref.get("filename") or ""), json.dumps(ref.get("title_path") or [], ensure_ascii=False), json.dumps(ref.get("keywords") or [], ensure_ascii=False), str(ref.get("brief") or ""), str(ref.get("content") or ref.get("_text") or ""), str(ref.get("chunk_type") or "text"), str(ref.get("content_hash") or ""), datetime.now(), ), ) async def upsert_wiki_page(page: Dict[str, Any]) -> int: async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: kb_id = str(page.get("kb_id") or "") slug = str(page.get("slug") or "")[:300] title = str(page.get("title") or "")[:300] page_type = str(page.get("page_type") or "topic") params = ( kb_id, slug, title, page_type, str(page.get("summary") or ""), str(page.get("content") or ""), json.dumps(page.get("keywords") or [], ensure_ascii=False), json.dumps(page.get("aliases") or [], ensure_ascii=False), json.dumps(page.get("related_file_ids") or [], ensure_ascii=False), json.dumps(page.get("source_refs") or [], ensure_ascii=False), json.dumps(page.get("chunk_refs") or [], ensure_ascii=False), str(page.get("source_kind") or "rule"), float(page.get("score_hint") or 1.0), datetime.now(), ) if slug: await cur.execute( """ UPDATE wiki_pages SET title = %s, page_type = %s, summary = %s, content = %s, keywords = %s::jsonb, aliases = %s::jsonb, related_file_ids = %s::jsonb, source_refs = %s::jsonb, chunk_refs = %s::jsonb, source_kind = %s, score_hint = %s, updated_at = %s WHERE kb_id = %s AND slug = %s RETURNING id """, ( title, page_type, str(page.get("summary") or ""), str(page.get("content") or ""), json.dumps(page.get("keywords") or [], ensure_ascii=False), json.dumps(page.get("aliases") or [], ensure_ascii=False), json.dumps(page.get("related_file_ids") or [], ensure_ascii=False), json.dumps(page.get("source_refs") or [], ensure_ascii=False), json.dumps(page.get("chunk_refs") or [], ensure_ascii=False), str(page.get("source_kind") or "rule"), float(page.get("score_hint") or 1.0), datetime.now(), kb_id, slug, ), ) row = await cur.fetchone() if row: return int(row[0]) await cur.execute( """ INSERT INTO wiki_pages ( kb_id, slug, title, page_type, summary, content, keywords, aliases, related_file_ids, source_refs, chunk_refs, source_kind, score_hint, created_at, updated_at ) VALUES (%s, %s, %s, %s, %s, %s, %s::jsonb, %s::jsonb, %s::jsonb, %s::jsonb, %s::jsonb, %s, %s, %s, %s) ON CONFLICT (kb_id, title, page_type) DO UPDATE SET slug = EXCLUDED.slug, summary = EXCLUDED.summary, content = EXCLUDED.content, keywords = EXCLUDED.keywords, aliases = EXCLUDED.aliases, related_file_ids = EXCLUDED.related_file_ids, source_refs = EXCLUDED.source_refs, chunk_refs = EXCLUDED.chunk_refs, source_kind = EXCLUDED.source_kind, score_hint = EXCLUDED.score_hint, updated_at = EXCLUDED.updated_at RETURNING id """, params + (datetime.now(),), ) row = await cur.fetchone() return int(row[0]) async def replace_page_slices(page_id: int, links: Iterable[Dict[str, Any]]) -> None: async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute("DELETE FROM wiki_page_slices WHERE wiki_page_id = %s", (page_id,)) for link in links: await cur.execute( """ INSERT INTO wiki_page_slices (wiki_page_id, file_id, slice_id, weight, reason) VALUES (%s, %s, %s, %s, %s) ON CONFLICT (wiki_page_id, slice_id) DO UPDATE SET weight = EXCLUDED.weight, reason = EXCLUDED.reason """, ( page_id, str(link.get("file_id") or ""), str(link.get("slice_id") or ""), float(link.get("weight") or 1.0), str(link.get("reason") or ""), ), ) async def acquire_wiki_page_lock(kb_id: str, slug: str, owner: str, ttl_seconds: int = 600) -> bool: kb_id = str(kb_id or "") slug = str(slug or "")[:300] owner = str(owner or "") ttl_seconds = max(int(ttl_seconds or 600), 30) now = datetime.now() expires_at = now + timedelta(seconds=ttl_seconds) async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( "DELETE FROM wiki_page_locks WHERE expires_at < %s", (now,), ) await cur.execute( """ INSERT INTO wiki_page_locks (kb_id, slug, owner, locked_at, expires_at) VALUES (%s, %s, %s, %s, %s) ON CONFLICT (kb_id, slug) DO NOTHING """, (kb_id, slug, owner, now, expires_at), ) return bool(cur.rowcount) async def release_wiki_page_lock(kb_id: str, slug: str, owner: str) -> None: async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( "DELETE FROM wiki_page_locks WHERE kb_id = %s AND slug = %s AND owner = %s", (str(kb_id or ""), str(slug or "")[:300], str(owner or "")), ) async def enqueue_wiki_finalize_op(job_id: str, kb_id: str, delay_seconds: int = 20) -> Optional[int]: kb_id = str(kb_id or "") if not kb_id: return None now = datetime.now() run_after = now + timedelta(seconds=max(int(delay_seconds or 0), 0)) payload = {"job_id": str(job_id or ""), "kb_id": kb_id} async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ UPDATE wiki_finalize_ops SET job_id = %s, payload = %s::jsonb, run_after = GREATEST(run_after, %s), updated_at = %s WHERE kb_id = %s AND status IN ('pending', 'running') RETURNING id """, (str(job_id or ""), json.dumps(payload, ensure_ascii=False), run_after, now, kb_id), ) row = await cur.fetchone() if row: return int(row[0]) await cur.execute( """ INSERT INTO wiki_finalize_ops ( job_id, kb_id, status, payload, claimed_at, run_after, enqueued_at, updated_at ) VALUES (%s, %s, 'pending', %s::jsonb, NULL, %s, %s, %s) RETURNING id """, (str(job_id or ""), kb_id, json.dumps(payload, ensure_ascii=False), run_after, now, now), ) row = await cur.fetchone() return int(row[0]) if row else None async def claim_wiki_finalize_ops(limit: int = 1) -> List[Dict[str, Any]]: limit = max(1, min(int(limit or 1), 10)) async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.transaction(): async with conn.cursor() as cur: await cur.execute( """ SELECT id FROM wiki_finalize_ops WHERE status = 'pending' AND run_after <= %s ORDER BY run_after ASC, id ASC LIMIT %s FOR UPDATE SKIP LOCKED """, (datetime.now(), limit), ) ids = [int(row[0]) for row in await cur.fetchall()] if not ids: return [] await cur.execute( """ UPDATE wiki_finalize_ops SET status = 'running', claimed_at = %s, updated_at = %s WHERE id = ANY(%s) RETURNING id, job_id, kb_id, payload """, (datetime.now(), datetime.now(), ids), ) rows = await cur.fetchall() keys = ["id", "job_id", "kb_id", "payload"] return [_row_to_dict(keys, row) for row in rows] async def complete_wiki_finalize_op(op_id: int) -> None: async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( "UPDATE wiki_finalize_ops SET status = 'done', updated_at = %s WHERE id = %s", (datetime.now(), int(op_id)), ) async def fail_wiki_finalize_op(op_id: int, error: str) -> None: payload = {"error": str(error or "")[:4000]} async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ UPDATE wiki_finalize_ops SET status = 'failed', payload = payload || %s::jsonb, updated_at = %s WHERE id = %s """, (json.dumps(payload, ensure_ascii=False), datetime.now(), int(op_id)), ) async def count_pages_for_kb(kb_id: str) -> int: async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute("SELECT COUNT(*) FROM wiki_pages WHERE kb_id = %s", (str(kb_id or ""),)) row = await cur.fetchone() return int(row[0] or 0) if row else 0 async def get_slice_refs_by_ids(slice_ids: Sequence[str]) -> List[Dict[str, Any]]: ids = [str(x) for x in slice_ids if str(x)] if not ids: return [] async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ SELECT kb_id, file_id, slice_id, ordinal, filename, title_path, keywords, brief, content, chunk_type, content_hash, updated_at FROM wiki_slice_refs WHERE slice_id = ANY(%s) """, (ids,), ) rows = await cur.fetchall() keys = [ "kb_id", "file_id", "slice_id", "ordinal", "filename", "title_path", "keywords", "brief", "content", "chunk_type", "content_hash", "updated_at", ] by_id = {str(item.get("slice_id")): item for item in (_row_to_dict(keys, row) for row in rows)} return [by_id[sid] for sid in ids if sid in by_id] async def search_slices_text(query: str, kb_id: Optional[str] = None, limit: int = 32) -> List[Dict[str, Any]]: terms = [t for t in _extract_search_terms(query) if t] if not terms: return [] where = [] params: List[Any] = [] if kb_id: where.append("sr.kb_id = %s") params.append(str(kb_id)) term_clauses = [] for term in terms[:8]: like = f"%{term}%" term_clauses.append( "(sr.filename ILIKE %s OR sr.title_path::text ILIKE %s OR sr.keywords::text ILIKE %s " "OR sr.brief ILIKE %s OR sr.content ILIKE %s)" ) params.extend([like, like, like, like, like]) where.append("(" + " OR ".join(term_clauses) + ")") params.append(max(int(limit or 32) * 4, int(limit or 32))) async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( f""" SELECT sr.kb_id, sr.file_id, sr.slice_id, sr.ordinal, sr.filename, sr.title_path, sr.keywords, sr.brief, sr.content, sr.chunk_type, sr.content_hash, sr.updated_at FROM wiki_slice_refs sr WHERE {" AND ".join(where)} ORDER BY sr.updated_at DESC LIMIT %s """, params, ) rows = await cur.fetchall() keys = [ "kb_id", "file_id", "slice_id", "ordinal", "filename", "title_path", "keywords", "brief", "content", "chunk_type", "content_hash", "updated_at", ] results = [] for row in rows: item = _row_to_dict(keys, row) item["score"] = _score_slice(query, item) results.append(item) results.sort(key=lambda x: float(x.get("score") or 0.0), reverse=True) return results[: max(int(limit or 32), 1)] async def search_pages(query: str, kb_id: Optional[str] = None, limit: int = 8) -> List[Dict[str, Any]]: terms = [t for t in _extract_search_terms(query) if t] where = [] params: List[Any] = [] rank_parts = [] rank_params: List[Any] = [] if kb_id: where.append("p.kb_id = %s") params.append(str(kb_id)) if terms: term_clauses = [] for term in terms[:6]: like = f"%{term}%" term_clauses.append( "(p.title ILIKE %s OR p.slug ILIKE %s OR p.summary ILIKE %s " "OR p.content ILIKE %s OR p.keywords::text ILIKE %s OR p.aliases::text ILIKE %s)" ) params.extend([like, like, like, like, like, like]) weight = _term_weight(str(term).lower()) rank_parts.append( "(CASE " "WHEN p.title ILIKE %s THEN %s " "WHEN p.slug ILIKE %s THEN %s " "WHEN p.aliases::text ILIKE %s THEN %s " "WHEN p.keywords::text ILIKE %s THEN %s " "WHEN p.summary ILIKE %s THEN %s " "WHEN p.content ILIKE %s THEN %s " "ELSE 0 END)" ) rank_params.extend([ like, 40.0 * weight, like, 25.0 * weight, like, 18.0 * weight, like, 12.0 * weight, like, 6.0 * weight, like, 1.0 * weight, ]) where.append("(" + " OR ".join(term_clauses) + ")") where_sql = " AND ".join(where) if where else "1=1" rank_sql = " + ".join(rank_parts) if rank_parts else "0" execute_params = rank_params + params + [max(limit * 10, limit, 80)] async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( f""" SELECT p.id, p.kb_id, p.slug, p.title, p.page_type, p.summary, p.content, p.keywords, p.aliases, p.related_file_ids, p.source_refs, p.chunk_refs, p.source_kind, p.score_hint, p.updated_at, ({rank_sql}) AS db_match_rank FROM wiki_pages p WHERE {where_sql} ORDER BY db_match_rank DESC, p.score_hint DESC, p.updated_at DESC LIMIT %s """, execute_params, ) rows = await cur.fetchall() keys = [ "id", "kb_id", "slug", "title", "page_type", "summary", "content", "keywords", "aliases", "related_file_ids", "source_refs", "chunk_refs", "source_kind", "score_hint", "updated_at", "db_match_rank", ] results = [] for row in rows: item = _row_to_dict(keys, row) item["score"] = _score_page(query, item) results.append(item) results.sort(key=lambda x: x.get("score", 0), reverse=True) return results[:limit] async def get_page(page_id: int) -> Optional[Dict[str, Any]]: async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ SELECT id, kb_id, slug, title, page_type, summary, content, keywords, aliases, related_file_ids, source_refs, chunk_refs, source_kind, score_hint, updated_at FROM wiki_pages WHERE id = %s """, (page_id,), ) row = await cur.fetchone() if not row: return None keys = [ "id", "kb_id", "slug", "title", "page_type", "summary", "content", "keywords", "aliases", "related_file_ids", "source_refs", "chunk_refs", "source_kind", "score_hint", "updated_at", ] page = _row_to_dict(keys, row) await cur.execute( """ SELECT ps.file_id, ps.slice_id, ps.weight, ps.reason, sr.filename, sr.title_path, sr.keywords, sr.brief, sr.content, sr.chunk_type FROM wiki_page_slices ps LEFT JOIN wiki_slice_refs sr ON sr.slice_id = ps.slice_id WHERE ps.wiki_page_id = %s ORDER BY ps.weight DESC LIMIT 50 """, (page_id,), ) links = await cur.fetchall() link_keys = ["file_id", "slice_id", "weight", "reason", "filename", "title_path", "keywords", "brief", "content", "chunk_type"] page["slices"] = [_row_to_dict(link_keys, row) for row in links] return page async def list_wiki_pages_for_kb(kb_id: str) -> List[Dict[str, Any]]: async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ SELECT id, kb_id, slug, title, page_type, summary, content, keywords, aliases, related_file_ids, source_refs, chunk_refs, source_kind, score_hint, updated_at FROM wiki_pages WHERE kb_id = %s ORDER BY updated_at DESC """, (str(kb_id or ""),), ) rows = await cur.fetchall() keys = [ "id", "kb_id", "slug", "title", "page_type", "summary", "content", "keywords", "aliases", "related_file_ids", "source_refs", "chunk_refs", "source_kind", "score_hint", "updated_at", ] pages = [_row_to_dict(keys, row) for row in rows] if not pages: return [] page_ids = [int(page["id"]) for page in pages if page.get("id")] placeholders = ",".join(["%s"] * len(page_ids)) await cur.execute( f""" SELECT ps.wiki_page_id, sr.kb_id, ps.file_id, ps.slice_id, ps.weight, ps.reason, sr.filename, sr.title_path, sr.keywords, sr.brief, sr.content, sr.chunk_type FROM wiki_page_slices ps LEFT JOIN wiki_slice_refs sr ON sr.slice_id = ps.slice_id WHERE ps.wiki_page_id IN ({placeholders}) ORDER BY ps.wiki_page_id, ps.weight DESC """, page_ids, ) link_rows = await cur.fetchall() link_keys = [ "wiki_page_id", "kb_id", "file_id", "slice_id", "weight", "reason", "filename", "title_path", "keywords", "brief", "content", "chunk_type", ] links_by_page: Dict[int, List[Dict[str, Any]]] = {} for row in link_rows: item = _row_to_dict(link_keys, row) page_id = int(item.pop("wiki_page_id")) links_by_page.setdefault(page_id, []).append(item) for page in pages: page["slices"] = links_by_page.get(int(page["id"]), []) return pages async def get_page_slices(page_ids: Sequence[int], max_slices_per_page: int = 8) -> Dict[int, List[Dict[str, Any]]]: if not page_ids: return {} placeholders = ",".join(["%s"] * len(page_ids)) async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( f""" SELECT ps.wiki_page_id, sr.kb_id, ps.file_id, ps.slice_id, ps.weight, ps.reason, sr.filename, sr.title_path, sr.keywords, sr.brief, sr.content, sr.chunk_type FROM wiki_page_slices ps LEFT JOIN wiki_slice_refs sr ON sr.slice_id = ps.slice_id WHERE ps.wiki_page_id IN ({placeholders}) ORDER BY ps.wiki_page_id, ps.weight DESC """, list(page_ids), ) rows = await cur.fetchall() out: Dict[int, List[Dict[str, Any]]] = {} keys = ["wiki_page_id", "kb_id", "file_id", "slice_id", "weight", "reason", "filename", "title_path", "keywords", "brief", "content", "chunk_type"] for row in rows: item = _row_to_dict(keys, row) page_id = int(item.pop("wiki_page_id")) bucket = out.setdefault(page_id, []) if len(bucket) < max_slices_per_page: bucket.append(item) return out async def log_query(query: str, kb_id: Optional[str], result_count: int) -> None: async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ INSERT INTO wiki_query_logs (query, kb_id, result_count, created_at) VALUES (%s, %s, %s, %s) """, (query, kb_id or "", result_count, datetime.now()), ) async def get_stats() -> Dict[str, Any]: async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: tables = [ "wiki_file_state", "wiki_slice_refs", "wiki_pages", "wiki_page_slices", "wiki_sync_jobs", "wiki_pending_ops", "wiki_dead_letters", "wiki_llm_cache", "wiki_page_locks", "wiki_finalize_ops", ] stats: Dict[str, Any] = {} for table in tables: await cur.execute(f"SELECT COUNT(*) FROM {table}") stats[table] = (await cur.fetchone())[0] if await _table_exists(cur, "wiki_slice_embeddings"): await cur.execute("SELECT COUNT(*) FROM wiki_slice_embeddings") stats["wiki_slice_embeddings"] = (await cur.fetchone())[0] return stats async def list_file_states(kb_ids: Sequence[str]) -> List[Dict[str, Any]]: ids = [str(x) for x in kb_ids if str(x)] if not ids: return [] async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ SELECT kb_id, file_id, filename, file_hash, status, slice_count, source_updated_at, metadata, last_synced_at FROM wiki_file_state WHERE kb_id = ANY(%s) """, (ids,), ) rows = await cur.fetchall() keys = [ "kb_id", "file_id", "filename", "file_hash", "status", "slice_count", "source_updated_at", "metadata", "last_synced_at", ] return [_row_to_dict(keys, row) for row in rows] async def delete_wiki_by_file(file_id: str) -> Dict[str, int]: file_id = str(file_id) counts = { "wiki_pages": 0, "wiki_page_slices": 0, "wiki_slice_embeddings": 0, "wiki_slice_refs": 0, "wiki_file_state": 0, } async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( "SELECT DISTINCT wiki_page_id FROM wiki_page_slices WHERE file_id = %s", (file_id,), ) page_ids = [row[0] for row in await cur.fetchall()] if page_ids: await cur.execute("DELETE FROM wiki_pages WHERE id = ANY(%s)", (page_ids,)) counts["wiki_pages"] = cur.rowcount or 0 counts["wiki_slice_embeddings"] = await _delete_optional_table_rows( cur, "wiki_slice_embeddings", "file_id = %s", [file_id] ) await cur.execute("DELETE FROM wiki_slice_refs WHERE file_id = %s", (file_id,)) counts["wiki_slice_refs"] = cur.rowcount or 0 await cur.execute("DELETE FROM wiki_file_state WHERE file_id = %s", (file_id,)) counts["wiki_file_state"] = cur.rowcount or 0 return counts async def delete_wiki_by_kb(kb_id: str) -> Dict[str, int]: kb_id = str(kb_id) counts = { "wiki_pages": 0, "wiki_slice_embeddings": 0, "wiki_slice_refs": 0, "wiki_file_state": 0, "wiki_sync_jobs": 0, "wiki_query_logs": 0, } async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute("DELETE FROM wiki_pages WHERE kb_id = %s", (kb_id,)) counts["wiki_pages"] = cur.rowcount or 0 counts["wiki_slice_embeddings"] = await _delete_optional_table_rows( cur, "wiki_slice_embeddings", "kb_id = %s", [kb_id] ) await cur.execute("DELETE FROM wiki_slice_refs WHERE kb_id = %s", (kb_id,)) counts["wiki_slice_refs"] = cur.rowcount or 0 await cur.execute("DELETE FROM wiki_file_state WHERE kb_id = %s", (kb_id,)) counts["wiki_file_state"] = cur.rowcount or 0 await cur.execute("DELETE FROM wiki_sync_jobs WHERE kb_id = %s", (kb_id,)) counts["wiki_sync_jobs"] = cur.rowcount or 0 await cur.execute("DELETE FROM wiki_query_logs WHERE kb_id = %s", (kb_id,)) counts["wiki_query_logs"] = cur.rowcount or 0 return counts async def delete_wiki_by_time( before: Optional[datetime] = None, after: Optional[datetime] = None, kb_id: Optional[str] = None, ) -> Dict[str, int]: if before is None and after is None: raise ValueError("before or after is required") counts = { "wiki_pages": 0, "wiki_slice_embeddings": 0, "wiki_slice_refs": 0, "wiki_file_state": 0, "wiki_sync_jobs": 0, "wiki_query_logs": 0, } where_pages, params_pages = _time_where("updated_at", before, after) where_slices, params_slices = _time_where("updated_at", before, after) where_files, params_files = _time_where("last_synced_at", before, after) where_jobs, params_jobs = _time_where("created_at", before, after) where_logs, params_logs = _time_where("created_at", before, after) if kb_id: where_pages += " AND kb_id = %s" params_pages.append(str(kb_id)) where_slices += " AND kb_id = %s" params_slices.append(str(kb_id)) where_files += " AND kb_id = %s" params_files.append(str(kb_id)) where_jobs += " AND kb_id = %s" params_jobs.append(str(kb_id)) where_logs += " AND kb_id = %s" params_logs.append(str(kb_id)) async with await _connect_pool() as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute(f"DELETE FROM wiki_pages WHERE {where_pages}", params_pages) counts["wiki_pages"] = cur.rowcount or 0 counts["wiki_slice_embeddings"] = await _delete_optional_table_rows( cur, "wiki_slice_embeddings", where_pages, params_pages ) await cur.execute(f"DELETE FROM wiki_slice_refs WHERE {where_slices}", params_slices) counts["wiki_slice_refs"] = cur.rowcount or 0 await cur.execute(f"DELETE FROM wiki_file_state WHERE {where_files}", params_files) counts["wiki_file_state"] = cur.rowcount or 0 await cur.execute(f"DELETE FROM wiki_sync_jobs WHERE {where_jobs}", params_jobs) counts["wiki_sync_jobs"] = cur.rowcount or 0 await cur.execute(f"DELETE FROM wiki_query_logs WHERE {where_logs}", params_logs) counts["wiki_query_logs"] = cur.rowcount or 0 return counts def _row_to_dict(keys: Sequence[str], row: Sequence[Any]) -> Dict[str, Any]: data = dict(zip(keys, row)) for key, value in list(data.items()): if isinstance(value, datetime): data[key] = value.isoformat() return data def _admin_plan_from_row(row: Optional[Sequence[Any]]) -> Dict[str, Any]: if not row: return {} keys = [ "id", "chat_id", "status", "intent", "scope", "kb_id", "kb_name", "summary", "actions", "executed", "metadata", "created_at", "updated_at", "expires_at", ] return _row_to_dict(keys, row) def _time_where( column: str, before: Optional[datetime], after: Optional[datetime], ) -> tuple[str, List[Any]]: parts = [] params: List[Any] = [] if before is not None: parts.append(f"{column} < %s") params.append(before) if after is not None: parts.append(f"{column} >= %s") params.append(after) return " AND ".join(parts) if parts else "1=0", params def _extract_search_terms(query: str) -> List[str]: import re query = (query or "").strip() if not query: return [] terms = [] for token in re.findall(r"[\u4e00-\u9fff]{2,12}|[A-Za-z0-9][A-Za-z0-9_\-]{1,30}", query): if token not in terms: terms.append(token) if re.fullmatch(r"[\u4e00-\u9fff]{3,12}", token): for size in (2, 3, 4, 5, 6): if len(token) < size: continue for idx in range(0, len(token) - size + 1): gram = token[idx:idx + size] if gram not in terms: terms.append(gram) compact_phrase = _compact_search_phrase(query) if compact_phrase and compact_phrase not in terms: terms.insert(0, compact_phrase) if query and query not in terms: terms.append(query[:80]) return terms[:40] _LOW_VALUE_SEARCH_TERMS = { "什么", "如何", "怎么", "哪些", "为什么", "是什么", "有哪些", "介绍", "说明", "查询", "检索", "请问", "帮我", "一下", "主要", "内容", "信息", "情况", "方法", "步骤", "这个", "那个", "是否", "有没有", "吗", "呢", "what", "how", "why", "which", "please", } def _compact_search_phrase(text: str) -> str: text = re.sub(r"\s+", "", (text or "").strip().lower()) for word in _LOW_VALUE_SEARCH_TERMS: text = text.replace(word, "") return text[:80] def _is_low_value_search_term(term: str) -> bool: term = (term or "").strip().lower() if not term: return True if term in _LOW_VALUE_SEARCH_TERMS: return True if re.fullmatch(r"\d+", term): return True if re.fullmatch(r"[\u4e00-\u9fff]", term): return True return False def _meaningful_search_terms(query: str) -> List[str]: terms = [] for term in _extract_search_terms(query): t = (term or "").strip().lower() if not t or _is_low_value_search_term(t): continue if t not in terms: terms.append(t) return terms[:24] def _score_page(query: str, page: Dict[str, Any]) -> float: title = str(page.get("title") or "").lower() slug = str(page.get("slug") or "").lower() summary = str(page.get("summary") or "").lower() content = str(page.get("content") or "")[:4000].lower() keywords = json.dumps(page.get("keywords") or [], ensure_ascii=False).lower() aliases = json.dumps(page.get("aliases") or [], ensure_ascii=False).lower() score = float(page.get("score_hint") or 1.0) compact_query = re.sub(r"\s+", "", query or "").lower() compact_phrase = _compact_search_phrase(compact_query) if compact_query: if compact_query == title: score += 40.0 elif compact_query in title: score += 25.0 elif compact_query in aliases: score += 18.0 elif compact_query in summary: score += 8.0 elif compact_query in content: score += 2.0 if compact_phrase and compact_phrase != compact_query: if compact_phrase == title: score += 36.0 elif compact_phrase in title: score += 24.0 elif compact_phrase in aliases: score += 18.0 elif compact_phrase in summary: score += 14.0 elif compact_phrase in content: score += 10.0 for term in _extract_search_terms(query): t = term.lower() if not t: continue weight = _term_weight(t) if title == t: score += 18.0 * weight elif slug == t: score += 14.0 * weight elif t in title: score += 9.0 * weight if t in aliases: score += 7.0 * weight if t in keywords: score += 4.0 * weight if t in summary: score += 2.0 * weight if t in slug: score += 2.0 * weight if t in content: score += 0.4 * weight meaningful_terms = _meaningful_search_terms(query) if meaningful_terms: title_text = f"{title} {aliases}" surface_text = f"{title} {aliases} {keywords} {summary}" full_text = f"{surface_text} {content}" matched_title = [t for t in meaningful_terms if t in title_text] matched_surface = [t for t in meaningful_terms if t in surface_text] matched_full = [t for t in meaningful_terms if t in full_text] score += 18.0 * (len(matched_title) / len(meaningful_terms)) score += 12.0 * (len(matched_surface) / len(meaningful_terms)) score += 18.0 * (len(matched_full) / len(meaningful_terms)) longest_query_term = max(meaningful_terms, key=len) if len(longest_query_term) >= 4 and longest_query_term not in full_text: score *= 0.72 specific_terms = [ (term.lower(), _term_weight(term.lower())) for term in _extract_search_terms(query) if _term_weight(term.lower()) >= 2.0 ] max_specific_weight = max((weight for _, weight in specific_terms), default=0.0) top_specific_terms = [ term for term, weight in specific_terms if weight >= max_specific_weight and term ] if ( str(page.get("page_type") or "") == "document_summary" and top_specific_terms and not any(term in title or term in aliases or term in slug for term in top_specific_terms) ): score *= 0.45 return score def _term_weight(term: str) -> float: if not term: return 0.0 if _is_low_value_search_term(term): return 0.2 if re.search(r"[a-z]", term) and re.search(r"[0-9]", term): return 3.0 if term in {"系统", "故障", "设备", "维修", "操作", "方法", "内容", "什么", "怎么", "如何", "哪些"}: return 0.2 if re.fullmatch(r"[\u4e00-\u9fff]{2}", term): return 1.2 if re.fullmatch(r"[\u4e00-\u9fff]+", term): return min(2.5, max(0.6, len(term) / 3.0)) return min(2.5, max(0.8, len(term) / 6.0)) def _score_slice(query: str, item: Dict[str, Any]) -> float: filename = str(item.get("filename") or "").lower() title = " ".join(str(x) for x in (item.get("title_path") or [])).lower() keywords = json.dumps(item.get("keywords") or [], ensure_ascii=False).lower() brief = str(item.get("brief") or "").lower() content = str(item.get("content") or "").lower() score = 0.0 for term in _extract_search_terms(query): t = term.lower() if not t: continue if t in filename: score += 3.0 if t in title: score += 2.5 if t in keywords: score += 2.0 if t in brief: score += 1.5 if t in content: score += 1.0 return score async def _table_exists(cur: Any, table_name: str) -> bool: try: await cur.execute("SELECT to_regclass(%s)", (table_name,)) row = await cur.fetchone() return bool(row and row[0]) except Exception: return False async def _delete_optional_table_rows( cur: Any, table_name: str, where_sql: str, params: Sequence[Any], ) -> int: if not await _table_exists(cur, table_name): return 0 await cur.execute(f"DELETE FROM {table_name} WHERE {where_sql}", list(params)) return cur.rowcount or 0