query.py

# -*- coding: utf-8 -*-
"""
query.py —— 群聊库检索接口（复用 build.py 的连接/向量逻辑；命令行独立入口）。

子命令:
  stats                  概览（消息/人数/时间跨度/Top发言人/类型/向量状态）
  who   [关键词]         发言人及发言量
  kw    "词"  [过滤]     关键词精确检索（LIKE 子串，短词也行）
  fts   "词组"[过滤]     全文检索（trigram，≥3字，bm25 排序）
  sem   "问题"[过滤]     语义检索（bge-m3 向量，找意思相近）
  ctx   <id> [-n 6]      某条消息前后上下文
  sql   "SELECT ..."     只读 SQL
  media <id> [-o 路径]   从 assets.db/files.db 取出原图/原文档
通用过滤(kw/fts/sem): --sender 名 --from 2026-01-01 --to 2026-03-01 --type 文本消息 -k 8
"""
import sys, os, argparse, sqlite3, datetime
import build

TMP_DIR = "tmp"


def con():
    return sqlite3.connect(build.DB)


def to_ts(s):
    return int(datetime.datetime.strptime(s, "%Y-%m-%d").timestamp())


def load_matrix(c):
    """读出全部向量 → (ids, mat (n,DIM))，供语义检索。"""
    import numpy as np
    ids = []; buf = bytearray()
    for mid, vec in c.execute("SELECT msg_id, vec FROM embeddings ORDER BY msg_id"):
        ids.append(mid); buf += vec
    if not ids:
        return np.empty(0, dtype=np.int64), np.empty((0, build.DIM), dtype=np.float32)
    return np.array(ids, dtype=np.int64), np.frombuffer(bytes(buf), dtype=np.float32).reshape(len(ids), build.DIM)


def filt_clause(a, prefix=""):
    p, cond, args = prefix, [], []
    if getattr(a, "sender", None):
        cond.append(f"({p}sender_name LIKE ? OR {p}sender_username = ?)"); args += [f"%{a.sender}%", a.sender]
    if getattr(a, "dfrom", None):
        cond.append(f"{p}create_time >= ?"); args.append(to_ts(a.dfrom))
    if getattr(a, "dto", None):
        cond.append(f"{p}create_time < ?"); args.append(to_ts(a.dto) + 86400)
    if getattr(a, "mtype", None):
        cond.append(f"{p}type = ?"); args.append(a.mtype)
    return cond, args


def show(rows, score=False):
    if not rows:
        print("（无结果）"); return
    for r in rows:
        head = f"#{r[0]} | {r[1]} | {r[2]}" + (f" | 相关度 {r[4]:.3f}" if score else "")
        body = (r[3] or "").strip().replace("\n", " ")
        if len(body) > 180: body = body[:180] + "…"
        print(f"[{head}]\n    {body}")


def cmd_stats(a):
    c = con(); q = lambda s: c.execute(s).fetchone()[0]
    print("会话    :", (build.get_state(c, "source_json", "?") or "?").split("/")[-1])
    print("消息总数:", q("SELECT COUNT(*) FROM messages"))
    print("发言人数:", q("SELECT COUNT(*) FROM senders"))
    print("时间跨度:", q("SELECT MIN(formatted_time) FROM messages"), "→", q("SELECT MAX(formatted_time) FROM messages"))
    print(f"已建向量: {build.get_state(c,'embed_count','0')} 条（模型 {build.get_state(c,'embed_model','无') or '无'}）")
    print("\n发言量 Top10:")
    for name, n in c.execute("SELECT last_name, msg_count FROM senders ORDER BY msg_count DESC LIMIT 10"):
        print(f"  {n:6d}  {name}")
    print("\n类型分布:")
    for t, n in c.execute("SELECT type, COUNT(*) FROM messages GROUP BY type ORDER BY 2 DESC"):
        print(f"  {n:6d}  {t}")


def cmd_who(a):
    c = con()
    if a.kw:
        rows = c.execute("SELECT last_name,username,msg_count FROM senders WHERE last_name LIKE ? ORDER BY msg_count DESC", (f"%{a.kw}%",)).fetchall()
    else:
        rows = c.execute("SELECT last_name,username,msg_count FROM senders ORDER BY msg_count DESC LIMIT 40").fetchall()
    for name, uid, n in rows:
        print(f"  {n:6d}  {name}   <{uid}>")


def cmd_kw(a):
    c = con(); cond, args = filt_clause(a)
    sql = "SELECT id,formatted_time,sender_name,content FROM messages WHERE " + " AND ".join(["content LIKE ?"] + cond) + " ORDER BY create_time LIMIT ?"
    show(c.execute(sql, [f"%{a.q}%"] + args + [a.limit]).fetchall())


def cmd_fts(a):
    c = con(); cond, args = filt_clause(a, prefix="m.")
    sql = ("SELECT m.id,m.formatted_time,m.sender_name,m.content,bm25(messages_fts) rk "
           "FROM messages_fts JOIN messages m ON m.id=messages_fts.rowid WHERE "
           + " AND ".join(["messages_fts MATCH ?"] + cond) + " ORDER BY rk LIMIT ?")
    try:
        rows = c.execute(sql, [a.q] + args + [a.limit]).fetchall()
    except sqlite3.OperationalError as e:
        print("FTS 出错（trigram 需≥3字，短词改用 kw）:", e); return
    show([(r[0], r[1], r[2], r[3]) for r in rows])


def cmd_sem(a):
    try:
        import numpy as np
    except ImportError:
        print("语义检索需 venv（numpy/torch），请用 .venv/bin/python 运行"); return
    c = con(); ids, mat = load_matrix(c)
    if len(ids) == 0:
        print("还没建向量，请先 .venv/bin/python build.py"); return
    qv = build.embed_texts([a.q], progress=False)[0]
    scores = mat @ qv
    cond, args = filt_clause(a)
    allow = {r[0] for r in c.execute("SELECT id FROM messages WHERE " + " AND ".join(cond), args)} if cond else None
    picked = []
    for idx in np.argsort(-scores):
        mid = int(ids[idx])
        if allow is not None and mid not in allow: continue
        picked.append((mid, float(scores[idx])))
        if len(picked) >= a.limit: break
    if not picked:
        show([]); return
    detail = {r[0]: r for r in c.execute(
        "SELECT id,formatted_time,sender_name,content FROM messages WHERE id IN (%s)" % ",".join(str(m) for m, _ in picked))}
    show([(*detail[m], s) for m, s in picked if m in detail], score=True)


def cmd_ctx(a):
    c = con()
    if not c.execute("SELECT 1 FROM messages WHERE id=?", (a.msg_id,)).fetchone():
        print("找不到该消息 id"); return
    for r in c.execute("SELECT id,formatted_time,sender_name,content FROM messages WHERE id BETWEEN ? AND ? ORDER BY id",
                       (a.msg_id - a.n, a.msg_id + a.n)).fetchall():
        mark = "→" if r[0] == a.msg_id else " "
        body = (r[3] or "").strip().replace("\n", " ")
        if len(body) > 180: body = body[:180] + "…"
        print(f"{mark} #{r[0]} | {r[1]} | {r[2]}: {body}")


def cmd_sql(a):
    c = con(); s = a.q.strip()
    if not s.lower().startswith(("select", "with", "explain")):
        print("只允许只读查询(SELECT/WITH/EXPLAIN)。"); return
    try:
        cur = c.execute(s)
        if cur.description: print(" | ".join(d[0] for d in cur.description))
        for r in cur.fetchall(): print(" | ".join("" if v is None else str(v) for v in r))
    except sqlite3.OperationalError as e:
        print("SQL 出错:", e)


def cmd_media(a):
    c = con(); base_dir = os.path.dirname(os.path.abspath(build.DB)); attached = []
    for dbfile, alias in [("assets.db", "ast"), ("files.db", "fil")]:
        p = os.path.join(base_dir, dbfile)
        if os.path.exists(p):
            c.execute(f"ATTACH ? AS {alias}", (p,)); attached.append(alias)
    row = None
    for alias in attached:
        try:
            row = c.execute(f"SELECT b.filename,s.blob FROM {alias}.media_blob b "
                            f"JOIN {alias}.blob_store s ON b.content_hash=s.content_hash "
                            f"WHERE b.msg_id=?", (a.msg_id,)).fetchone()
            if row: break
        except sqlite3.OperationalError:
            pass
    if not row:
        print(f"消息 #{a.msg_id} 没有已入库的媒体。"); return
    filename, blob = row
    if not a.out:
        os.makedirs(TMP_DIR, exist_ok=True)
    out = a.out or os.path.join(TMP_DIR, filename or f"media_{a.msg_id}.bin")
    with open(out, "wb") as f: f.write(blob)
    print(f"已导出: {out}  ({len(blob):,} 字节)")


def main():
    ap = argparse.ArgumentParser(description="群聊库检索")
    sub = ap.add_subparsers(dest="cmd", required=True)
    def add_filters(p):
        p.add_argument("--sender"); p.add_argument("--from", dest="dfrom")
        p.add_argument("--to", dest="dto"); p.add_argument("--type", dest="mtype")
        p.add_argument("-k", "--limit", type=int, default=8)
    sub.add_parser("stats")
    pw = sub.add_parser("who"); pw.add_argument("kw", nargs="?")
    for name in ("kw", "fts", "sem"):
        p = sub.add_parser(name); p.add_argument("q"); add_filters(p)
    pc = sub.add_parser("ctx"); pc.add_argument("msg_id", type=int); pc.add_argument("-n", type=int, default=6)
    ps = sub.add_parser("sql"); ps.add_argument("q")
    pm = sub.add_parser("media"); pm.add_argument("msg_id", type=int); pm.add_argument("-o", "--out")
    a = ap.parse_args()
    {"stats": cmd_stats, "who": cmd_who, "kw": cmd_kw, "fts": cmd_fts, "sem": cmd_sem,
     "ctx": cmd_ctx, "sql": cmd_sql, "media": cmd_media}[a.cmd](a)


if __name__ == "__main__":
    main()