hello-agents/Co-creation-projects/DeLunnLi-PaperGraph/backend/app/services/graph/kg_relations.py

"""知识图谱关系管理 —— 节点/边数据模型、图谱指标统计与查询."""

from __future__ import annotations

import json
import logging
import re
import sqlite3
import threading
import time
from typing import Any

from ...utils.common import suppress_exceptions

from ..reader.paper_reader_context import extract_pdf_text_full_cached, extract_pdf_text_full, _cache_set
from ...utils import normalize_arxiv_id as _norm_arxiv_id

logger = logging.getLogger(__name__)

_kg_infer_lock = threading.Lock()

_kg_recent_fingerprints: dict[str, float] = {}
_kg_fingerprints_lock = threading.Lock()
_KG_DEDUP_WINDOW_SEC = 180.0

_kg_metrics: dict[str, int] = {
    "build_ok": 0,
    "build_skip_no_candidates": 0,
    "build_skip_dedup": 0,
    "relations_upserted": 0,
}
_kg_metrics_lock = threading.Lock()

def get_kg_metrics() -> dict[str, int]:
    with _kg_metrics_lock:
        return dict(_kg_metrics)

def _prune_recent_fingerprints(now: float) -> None:
    cutoff = now - _KG_DEDUP_WINDOW_SEC * 2
    with _kg_fingerprints_lock:
        dead = [k for k, t in _kg_recent_fingerprints.items() if t < cutoff]
        for k in dead:
            _kg_recent_fingerprints.pop(k, None)

from ...utils.common import exec_sql

def ensure_tables(db_path: str) -> None:
    exec_sql(db_path,
        """CREATE TABLE IF NOT EXISTS paper_relations (
          source_paper_id INTEGER NOT NULL,
          target_paper_id INTEGER NOT NULL,
          relation TEXT NOT NULL,
          score REAL DEFAULT 0.0,
          evidence TEXT,
          created_at INTEGER NOT NULL,
          updated_at INTEGER NOT NULL,
          PRIMARY KEY (source_paper_id, target_paper_id, relation)
        )""",
        "CREATE INDEX IF NOT EXISTS idx_paper_relations_source ON paper_relations(source_paper_id)",
        "CREATE INDEX IF NOT EXISTS idx_paper_relations_target ON paper_relations(target_paper_id)",
    )

@suppress_exceptions(default_return=[])
def _loads_json_list(x: str | None) -> list[Any]:
    r = json.loads(x or "[]")
    return r if isinstance(r, list) else []

def _row_to_paper_meta(row: sqlite3.Row) -> dict[str, Any]:
    keys = row.keys()

    def _col(name: str) -> Any:
        return row[name] if name in keys else None

    return {
        "id": row["id"],
        "title": row["title"],
        "abstract": row["abstract"],
        "keywords": _loads_json_list(_col("keywords")),
        "tags": _loads_json_list(_col("tags")),
        "journal": row["journal"],
        "venue_type": row["venue_type"] if "venue_type" in keys else None,
        "year": row["year"],
        "category": row["category"],
        "local_pdf_path": _col("local_pdf_path"),
        "doi": (_col("doi") or None),
        "arxiv_id": (_col("arxiv_id") or None),
        "references": _loads_json_list(_col("references")),
    }

def _pdf_abspath_from_row(db_path: str, local_pdf_path: str | None) -> str | None:
    import os

    if not local_pdf_path or not str(local_pdf_path).strip():
        return None
    data_root = os.path.dirname(os.path.abspath(db_path))
    abspath = os.path.normpath(os.path.join(data_root, str(local_pdf_path).strip()))
    if os.path.isfile(abspath):
        return abspath
    return None

def _extract_related_work_excerpt(text: str, max_chars: int = 1600) -> str:
    t = (text or "").strip()
    if not t:
        return ""
    patterns = [
        r"(?i)\brelated\s+work\b",
        r"(?i)\brelated\s+works\b",
        r"相关工作",
    ]
    for pat in patterns:
        m = re.search(pat, t)
        if m:
            return t[m.start() : m.start() + max_chars].strip()
    return t[: max_chars // 2].strip()

_TOKEN_RE = re.compile(r"[^\w\u4e00-\u9fff]+", re.UNICODE)

_LEX_STOP = frozenset({
    "the", "a", "an", "and", "or", "of", "to", "in", "for", "with", "on", "by",
    "we", "our", "is", "are", "be", "this", "that", "from", "as", "at", "it",
})

def _lexical_token_set(meta: dict[str, Any]) -> set[str]:
    parts = [
        str(meta.get("title") or ""),
        str(meta.get("abstract") or ""),
        *(str(k) for k in meta.get("keywords") or []),
        *(str(t) for t in meta.get("tags") or []),
    ]
    text = " ".join(parts).lower()
    text = _TOKEN_RE.sub(" ", text)
    return {s for s in text.split() if len(s) >= 2 and s not in _LEX_STOP}

def _reference_signatures(refs: list[Any]) -> set[str]:
    sigs: set[str] = set()
    for r in refs or []:
        s = str(r).strip().lower()
        if not s:
            continue
        sigs.add(s)
        m = re.search(r"(10\.\d{4,9}/[^\s,;\"'<>]+)", s)
        if m:
            sigs.add(m.group(1).rstrip(").,]}"))
        m = re.search(r"arxiv[:/\s]*(\d{4}\.\d{4,5})(?:v\d+)?", s)
        if m:
            sigs.add(_norm_arxiv_id(m.group(1)) or m.group(1))
    return sigs

def _ref_overlap_bonus(new_sigs: set[str], cand: dict[str, Any]) -> float:
    if not new_sigs:
        return 0.0
    doi = (cand.get("doi") or "").strip().lower()
    if doi and doi in new_sigs:
        return 1.0
    ax = _norm_arxiv_id(cand.get("arxiv_id"))
    if ax and ax in new_sigs:
        return 1.0
    if doi:
        for sig in new_sigs:
            if len(sig) > 8 and sig in doi:
                return 0.85
    return 0.0

def fetch_new_and_candidates(
    db_path: str,
    new_paper_id: int,
    k: int = 32,
    *,
    sql_limit: int = 22,
    lexical_pool: int = 480,
) -> tuple[dict[str, Any | None], list[dict[str, Any]]]:
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    cur = conn.cursor()
    cur.execute("SELECT * FROM papers WHERE id=?", (int(new_paper_id),))
    row = cur.fetchone()
    if not row:
        conn.close()
        return None, []
    new_meta = _row_to_paper_meta(row)

    cat = (new_meta.get("category") or "").strip()
    year = new_meta.get("year")
    params: list[Any] = []
    where = ["id != ?"]
    params.append(int(new_paper_id))

    if cat:
        where.append("(category = ? OR category LIKE ?)")
        params.extend([cat, f"{cat.split('/')[0]}%"])

    if year:
        try:
            y = int(year)
            where.append("(year IS NULL OR year >= ?)")
            params.append(max(1900, y - 8))
        except Exception:
            pass

    wsql = " AND ".join(where)
    cur.execute(
        f"SELECT * FROM papers WHERE {wsql} ORDER BY created_at DESC LIMIT ?",
        (*params, int(sql_limit)),
    )
    sql_metas = [_row_to_paper_meta(r) for r in cur.fetchall()]

    cur.execute(
        """
        SELECT * FROM papers
        WHERE id != ?
        ORDER BY created_at DESC
        LIMIT ?
        """,
        (int(new_paper_id), int(lexical_pool)),
    )
    pool_rows = cur.fetchall()
    conn.close()

    new_lex = _lexical_token_set(new_meta)
    new_ref_sigs = _reference_signatures(new_meta.get("references") or [])

    scored: list[tuple[float, dict[str, Any]]] = []
    for r in pool_rows:
        meta = _row_to_paper_meta(r)
        cand_tokens = _lexical_token_set(meta)
        j = len(new_lex & cand_tokens) / max(1, len(new_lex | cand_tokens))
        ro = _ref_overlap_bonus(new_ref_sigs, meta)
        comb = j + ro * 0.45
        scored.append((comb, meta))

    scored.sort(key=lambda x: -x[0])

    seen: set[int] = set()
    out: list[dict[str, Any]] = []

    for m in sql_metas:
        pid = int(m["id"])
        if pid not in seen:
            seen.add(pid)
            out.append(m)

    min_lex = 0.055
    for comb, m in scored:
        if len(out) >= int(k):
            break
        pid = int(m["id"])
        if pid in seen:
            continue
        ro = _ref_overlap_bonus(new_ref_sigs, m)
        if comb < min_lex and ro <= 0:
            continue
        seen.add(pid)
        out.append(m)

    return new_meta, out[: int(k)]

def upsert_relations(db_path: str, source_paper_id: int, edges: list[dict[str, Any]]) -> int:
    ensure_tables(db_path)
    now = int(time.time())
    conn = sqlite3.connect(db_path)
    cur = conn.cursor()
    n = 0
    for e in edges:
        try:
            tid = int(e.get("target_paper_id"))
        except Exception:
            continue
        rel = str(e.get("relation") or "").strip()[:32]
        if len(rel) < 2 or len(rel) > 32:
            continue
        try:
            score = float(e.get("score") or 0.0)
        except Exception:
            score = 0.0
        ev = str(e.get("evidence") or "").strip()[:240]
        if tid <= 0 or tid == int(source_paper_id):
            continue
        cur.execute(
            """
            INSERT INTO paper_relations(source_paper_id, target_paper_id, relation, score, evidence, created_at, updated_at)
            VALUES(?, ?, ?, ?, ?, ?, ?)
            ON CONFLICT(source_paper_id, target_paper_id, relation) DO UPDATE SET
              score=excluded.score,
              evidence=excluded.evidence,
              updated_at=excluded.updated_at
            """,
            (int(source_paper_id), int(tid), rel, float(score), ev, now, now),
        )
        n += 1
    conn.commit()
    conn.close()
    return n

def build_relations_for_new_paper(db_path: str, new_paper_id: int) -> int:
    from ...agents import get_knowledge_graph_agent

    ensure_tables(db_path)
    new_meta, cands = fetch_new_and_candidates(db_path, int(new_paper_id), k=32)
    if not new_meta:
        return 0
    if not cands:
        with _kg_metrics_lock:
            _kg_metrics["build_skip_no_candidates"] = _kg_metrics.get("build_skip_no_candidates", 0) + 1
        logger.info(
            "kg_build_skip_no_candidates",
            extra={"paper_id": int(new_paper_id)},
        )
        return 0

    now = time.time()
    _prune_recent_fingerprints(now)

    _ax = _norm_arxiv_id(new_meta.get("arxiv_id"))
    if _ax:
        fp = f"arxiv:{_ax}"
    else:
        _doi = (new_meta.get("doi") or "").strip().lower()
        if _doi:
            fp = f"doi:{_doi}"
        else:
            _t = (new_meta.get("title") or "").strip().lower()[:160]
            fp = f"title:{_t}"

    with _kg_fingerprints_lock:
        last = _kg_recent_fingerprints.get(fp)
        if last is not None and (now - last) < _KG_DEDUP_WINDOW_SEC:
            with _kg_metrics_lock:
                _kg_metrics["build_skip_dedup"] = _kg_metrics.get("build_skip_dedup", 0) + 1
            logger.info(
                "kg_build_skip_dedup",
                extra={"paper_id": int(new_paper_id), "fingerprint": fp},
            )
            return 0
        _kg_recent_fingerprints[fp] = now

    try:
        pdf_abspath = _pdf_abspath_from_row(db_path, new_meta.get("local_pdf_path"))
        if pdf_abspath:
            excerpt, _hit = extract_pdf_text_full_cached(
                db_path, int(new_paper_id), pdf_abspath, max_chars=9000
            )
            if not excerpt.strip():
                excerpt = extract_pdf_text_full(pdf_abspath, max_chars=9000)
                if excerpt.strip():
                    _cache_set(db_path, int(new_paper_id), pdf_abspath, excerpt)
            if excerpt.strip():
                new_meta["pdf_excerpt"] = excerpt[:9000]
                new_meta["related_work_excerpt"] = _extract_related_work_excerpt(excerpt, max_chars=1600)
    except Exception as exc:
        logger.warning(
            "kg_pdf_excerpt_failed",
            extra={"paper_id": int(new_paper_id)},
            exc_info=exc,
        )

    edges: list[dict[str, Any]] = []
    with _kg_infer_lock:
        agent = get_knowledge_graph_agent()
        edges, _ = agent.infer_edges(new_paper=new_meta, candidates=cands)

    n = upsert_relations(db_path, int(new_paper_id), edges)
    if n:
        with _kg_metrics_lock:
            _kg_metrics["relations_upserted"] = _kg_metrics.get("relations_upserted", 0) + n
    with _kg_metrics_lock:
        _kg_metrics["build_ok"] = _kg_metrics.get("build_ok", 0) + 1
    return n