1799 lines
78 KiB
Python
1799 lines
78 KiB
Python
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
|