sciwrite-lint 0.2.0__py3-none-any.whl

sciwrite_lint/references/crossref.py

@@ -0,0 +1,282 @@
+"""CrossRef API client — async httpx.
+
+CrossRef is the gold standard for DOI metadata. Provides authoritative
+title, authors, year, venue, and retraction status.
+
+Uses httpx.AsyncClient (shared with all other API clients), eliminating
+the per-call ~100-200 MB overhead that habanero/requests caused via
+separate SSL contexts and connection pools per thread.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Any
+
+import httpx
+from loguru import logger
+
+from sciwrite_lint._network import clean_and_validate_doi
+from sciwrite_lint.models import Citation
+
+_CROSSREF_BASE = "https://api.crossref.org/works"
+
+# Fields needed by _parse_crossref — avoids downloading full records
+# (reference lists alone can be 5-20 MB per paper).
+_CROSSREF_SELECT = (
+    "DOI,title,author,published-print,published-online,issued,"
+    "container-title,is-referenced-by-count,type,abstract,update-to"
+)
+
+
+def _crossref_year(item: dict[str, Any]) -> int | None:
+    """Extract publication year from a raw CrossRef work item."""
+    for date_field in ("published-print", "published-online", "issued"):
+        parts = item.get(date_field, {}).get("date-parts", [[]])
+        if parts and parts[0] and parts[0][0]:
+            return parts[0][0]
+    return None
+
+
+def _normalize_crossref_item(item: dict[str, Any]) -> dict[str, Any]:
+    """Convert raw CrossRef item to the normalized format _best_match expects."""
+    titles = item.get("title", [])
+    authors = []
+    for a in item.get("author", []):
+        given = a.get("given", "")
+        family = a.get("family", "")
+        if family:
+            authors.append(f"{given} {family}".strip() if given else family)
+    containers = item.get("container-title", [])
+    return {
+        "title": titles[0] if titles else "",
+        "authors": authors,
+        "year": _crossref_year(item),
+        "venue": containers[0] if containers else "",
+    }
+
+
+def _best_title_match(
+    query_title: str,
+    items: list[dict[str, Any]],
+    query_authors: list[str] | tuple[str, ...] | None = None,
+    query_year: int | None = None,
+    query_venue: str = "",
+) -> dict[str, Any] | None:
+    """Pick the CrossRef result whose title best matches the query.
+
+    Delegates scoring to ``_best_match`` (single implementation of
+    title + author + year + venue scoring). Returns the raw CrossRef item.
+    """
+    from sciwrite_lint.api import _best_match
+
+    # Build normalized candidates, keeping a map back to the raw items
+    normalized = []
+    raw_by_idx: dict[int, dict[str, Any]] = {}
+    for i, item in enumerate(items):
+        norm = _normalize_crossref_item(item)
+        norm["_idx"] = i
+        normalized.append(norm)
+        raw_by_idx[i] = item
+
+    best = _best_match(
+        query_title,
+        normalized,
+        query_authors=query_authors,
+        query_year=query_year,
+        query_venue=query_venue,
+    )
+    if best is None:
+        return None
+    return raw_by_idx[best["_idx"]]
+
+
+async def crossref_lookup(
+    citation: Citation,
+    polite_email: str = "",
+    client: httpx.AsyncClient | None = None,
+) -> dict[str, Any] | None:
+    """Look up a citation via CrossRef. Returns normalized dict or None.
+
+    Async — uses httpx. Caller should pass a shared client for connection
+    reuse. Rate limiting is the caller's responsibility.
+    """
+    _client = client or httpx.AsyncClient(timeout=15.0, follow_redirects=True)
+
+    headers = (
+        {"User-Agent": f"sciwrite-lint/1.0 (mailto:{polite_email})"}
+        if polite_email
+        else {}
+    )
+
+    try:
+        return await _do_crossref_lookup(citation, _client, headers, polite_email)
+    except Exception as e:
+        logger.debug("CrossRef lookup failed for {}: {}", citation.key, e)
+        return None
+    finally:
+        if client is None:
+            await _client.aclose()
+
+
+async def _do_crossref_lookup(
+    citation: Citation,
+    client: httpx.AsyncClient,
+    headers: dict[str, str],
+    polite_email: str = "",
+) -> dict[str, Any] | None:
+    """Core lookup logic with retries handled by retry_on_transient."""
+    from sciwrite_lint.rate_limiter import retry_on_transient
+
+    # mailto query param ensures polite pool (higher rate limits)
+    base_params: dict[str, str] = {}
+    if polite_email:
+        base_params["mailto"] = polite_email
+
+    # DOI lookup first (most reliable)
+    # Note: /works/{doi} does not support select= (returns 400)
+    clean_doi = clean_and_validate_doi(citation.doi) if citation.doi else None
+    if clean_doi:
+        resp = await retry_on_transient(
+            lambda: client.get(
+                f"{_CROSSREF_BASE}/{clean_doi}",
+                params=base_params,
+                headers=headers,
+            ),
+            label=f"CrossRef DOI {citation.key}",
+        )
+        if resp.status_code == 200:
+            msg = resp.json().get("message")
+            if msg:
+                return _parse_crossref(msg)
+
+    # Title search (DOI not available or failed)
+    if citation.title:
+        from sciwrite_lint.api import _first_surname
+
+        surname = _first_surname(citation.authors)
+        query = f"{surname} {citation.title}" if surname else citation.title
+        resp = await retry_on_transient(
+            lambda query=query: client.get(
+                _CROSSREF_BASE,
+                params={
+                    **base_params,
+                    "query": query,
+                    "rows": 10,
+                    "select": _CROSSREF_SELECT,
+                },
+                headers=headers,
+            ),
+            label=f"CrossRef title {citation.key}",
+        )
+        if resp.status_code == 200:
+            items = resp.json().get("message", {}).get("items", [])
+            if items:
+                best = _best_title_match(
+                    citation.title,
+                    items,
+                    query_authors=citation.authors,
+                    query_year=int(citation.year) if citation.year.isdigit() else None,
+                    query_venue=citation.venue,
+                )
+                if best is not None:
+                    return _parse_crossref(best)
+
+    return None
+
+
+async def check_retraction(
+    doi: str,
+    polite_email: str = "",
+    client: httpx.AsyncClient | None = None,
+) -> bool:
+    """Check if a DOI points to a retracted paper via CrossRef metadata."""
+    if not doi:
+        return False
+
+    _client = client or httpx.AsyncClient(timeout=15.0, follow_redirects=True)
+
+    headers = (
+        {"User-Agent": f"sciwrite-lint/1.0 (mailto:{polite_email})"}
+        if polite_email
+        else {}
+    )
+
+    try:
+        # /works/{doi} does not support select= (returns 400)
+        resp = await _client.get(
+            f"{_CROSSREF_BASE}/{doi}",
+            headers=headers,
+        )
+        if resp.status_code != 200:
+            return False
+        msg = resp.json().get("message", {})
+        updates = msg.get("update-to", [])
+        for update in updates:
+            if update.get("type") == "retraction":
+                return True
+            label = update.get("label", "").lower()
+            if "retract" in label:
+                return True
+        return False
+    except Exception as e:
+        logger.debug("Retraction check failed for DOI {}: {}", doi, e)
+        return False
+    finally:
+        if client is None:
+            await _client.aclose()
+
+
+def _parse_crossref(work: dict) -> dict[str, Any]:
+    """Parse CrossRef work into normalized dict."""
+    # Authors
+    authors = []
+    for a in work.get("author", []):
+        given = a.get("given", "")
+        family = a.get("family", "")
+        if family:
+            name = f"{given} {family}".strip() if given else family
+            authors.append(name)
+
+    # Title
+    titles = work.get("title", [])
+    title = titles[0] if titles else ""
+
+    # Year — prefer print date, then online
+    year = None
+    for date_field in ("published-print", "published-online", "issued"):
+        parts = work.get(date_field, {}).get("date-parts", [[]])
+        if parts and parts[0] and parts[0][0]:
+            year = parts[0][0]
+            break
+
+    # Venue
+    containers = work.get("container-title", [])
+    venue = containers[0] if containers else ""
+
+    # DOI
+    doi = work.get("DOI", "")
+
+    # Abstract
+    abstract = work.get("abstract", "")
+    if abstract:
+        abstract = re.sub(r"<[^>]+>", "", abstract).strip()
+
+    # Retraction
+    retracted = False
+    for update in work.get("update-to", []):
+        if "retract" in (update.get("type", "") + update.get("label", "")).lower():
+            retracted = True
+            break
+
+    return {
+        "source": "crossref",
+        "title": title,
+        "authors": authors,
+        "year": year,
+        "venue": venue,
+        "doi": doi,
+        "citation_count": work.get("is-referenced-by-count", 0),
+        "abstract": abstract or None,
+        "retracted": retracted,
+    }

sciwrite_lint/references/embedding_store.py

@@ -0,0 +1,380 @@
+"""SQLite-vec embedding storage for claim verification.
+
+Stores chunk text + embeddings in the per-paper workspace DB
+(``references/{paper}/parsed/workspace.db``) using sqlite-vec for KNN
+retrieval.
+
+OOM-safe: encodes and inserts in small batches (EMBED_BATCH_SIZE),
+never holding all vectors in RAM at once.
+"""
+
+from __future__ import annotations
+
+import sqlite3
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import numpy as np
+
+from loguru import logger
+
+from sciwrite_lint.references.workspace_db import (
+    db_path as _db_path,
+    open_db as _open_db,
+    serialize_f32 as _serialize_f32,
+)
+
+EMBED_BATCH_SIZE = 32  # chunks per encode + insert batch
+
+# Adaptive encode batch sizing based on token length.
+# Eager attention memory scales as O(batch * seq_len^2). With 20 GB VRAM
+# on RTX 4000 Ada and vLLM occupying ~18.6 GB (paged out on WSL2):
+#   ≤512 tokens, batch=32: ~2.5 GB peak — fits easily
+#   ≤2048 tokens, batch=8: ~5.5 GB peak — safe
+#   >2048 tokens, batch=4:  ~11 GB peak — safe on WSL2 (worst real: 3810 tokens)
+_ENCODE_TIERS: list[tuple[int, int]] = [
+    (512, 32),  # short chunks: full batch
+    (2048, 8),  # medium chunks: reduced batch
+]
+_ENCODE_BATCH_LONG = 4  # anything above the last tier threshold
+
+
+def _encode_adaptive(model: object, texts: list[str]) -> np.ndarray:
+    """Encode texts with batch size adapted to token length.
+
+    Short texts (≤512 tokens) are encoded in large batches for throughput.
+    Long texts get smaller batches to stay within GPU memory limits.
+    Results are reassembled in the original order.
+    """
+    import numpy as np
+
+    tokenizer = model.tokenizer  # type: ignore[attr-defined]
+    # Tokenize to get lengths (fast, CPU-only)
+    token_counts = [len(tokenizer(t, truncation=False)["input_ids"]) for t in texts]
+
+    # Group indices by tier
+    groups: dict[int, list[int]] = {}
+    for idx, tok_len in enumerate(token_counts):
+        batch_size = _ENCODE_BATCH_LONG
+        for threshold, bs in _ENCODE_TIERS:
+            if tok_len <= threshold:
+                batch_size = bs
+                break
+        groups.setdefault(batch_size, []).append(idx)
+
+    # Encode each group with its batch size, collect results
+    results: dict[int, np.ndarray] = {}
+    for batch_size, indices in groups.items():
+        group_texts = [texts[i] for i in indices]
+        vecs = model.encode(  # type: ignore[attr-defined]
+            group_texts, normalize_embeddings=True, batch_size=batch_size
+        )
+        for i, idx in enumerate(indices):
+            results[idx] = vecs[i]
+
+    # Reassemble in original order
+    return np.stack([results[i] for i in range(len(texts))])
+
+
+def _ensure_vec_table(conn: sqlite3.Connection, dim: int) -> None:
+    """Create the vec0 virtual table if it doesn't exist, or recreate on dimension change."""
+    # Check if table exists
+    row = conn.execute(
+        "SELECT name FROM sqlite_master WHERE type='table' AND name='chunk_embeddings'"
+    ).fetchone()
+
+    if row:
+        # Check dimension matches
+        stored_dim = conn.execute(
+            "SELECT value FROM embed_meta WHERE key='dimension'"
+        ).fetchone()
+        if stored_dim and int(stored_dim[0]) == dim:
+            return
+        # Dimension changed — drop and recreate
+        logger.info("Embedding dimension changed to {}, rebuilding vec table", dim)
+        conn.execute("DROP TABLE chunk_embeddings")
+
+    conn.execute(
+        f"CREATE VIRTUAL TABLE chunk_embeddings USING vec0("
+        f"embedding float[{dim}] distance_metric=cosine)"
+    )
+    conn.execute(
+        "INSERT OR REPLACE INTO embed_meta (key, value) VALUES ('dimension', ?)",
+        (str(dim),),
+    )
+    conn.commit()
+
+
+def has_embeddings(
+    ref_key: str, references_dir: Path, model_name: str | None = None
+) -> bool:
+    """Check if a reference has complete, current embeddings.
+
+    Returns False if:
+    - No embeddings exist for this key
+    - Embedding was interrupted (complete=0) — cleans up broken data
+    - Embedding model changed (stored model != current model)
+    """
+    db_file = _db_path(references_dir)
+    if not db_file.exists():
+        return False
+    try:
+        conn = _open_db(references_dir)
+
+        # Check completion status
+        row = conn.execute(
+            "SELECT complete FROM ref_status WHERE ref_key = ?", (ref_key,)
+        ).fetchone()
+        if row is None:
+            conn.close()
+            return False
+        if row[0] != 1:
+            # Incomplete — clean up broken data
+            logger.debug("Cleaning up incomplete embeddings for {}", ref_key)
+            _delete_ref_data(conn, ref_key)
+            conn.commit()
+            conn.close()
+            return False
+
+        # Check model compatibility
+        if model_name is not None:
+            stored = conn.execute(
+                "SELECT value FROM embed_meta WHERE key='model'"
+            ).fetchone()
+            if stored and stored[0] != model_name:
+                conn.close()
+                return False
+
+        conn.close()
+        return True
+    except Exception:
+        return False
+
+
+def _delete_ref_data(conn: sqlite3.Connection, ref_key: str) -> int:
+    """Delete all data for a reference (chunks, vectors, status). Returns count."""
+    existing_ids = conn.execute(
+        "SELECT id FROM chunks WHERE ref_key = ?", (ref_key,)
+    ).fetchall()
+    for (chunk_id,) in existing_ids:
+        conn.execute("DELETE FROM chunk_embeddings WHERE rowid = ?", (chunk_id,))
+    conn.execute("DELETE FROM chunks WHERE ref_key = ?", (ref_key,))
+    conn.execute("DELETE FROM ref_status WHERE ref_key = ?", (ref_key,))
+    return len(existing_ids)
+
+
+def store_embeddings(
+    ref_key: str,
+    chunks: list[dict],
+    references_dir: Path,
+    model_name: str,
+    embedding_dim: int,
+) -> int:
+    """Store chunk text + embeddings in batches (OOM-safe).
+
+    Each chunk dict must have: text, section_title, granularity, start_char.
+    Embeddings are computed internally in batches of EMBED_BATCH_SIZE.
+
+    Returns the number of chunks stored.
+    """
+    from sciwrite_lint.references.reference_store import _get_embedding_model
+
+    conn = _open_db(references_dir)
+    _ensure_vec_table(conn, embedding_dim)
+
+    # Clear existing data for this key (including incomplete runs)
+    _delete_ref_data(conn, ref_key)
+
+    # Mark as in-progress (complete=0) with expected count
+    conn.execute(
+        "INSERT INTO ref_status (ref_key, expected_chunks, stored_chunks, complete) "
+        "VALUES (?, ?, 0, 0)",
+        (ref_key, len(chunks)),
+    )
+    conn.commit()
+
+    model = _get_embedding_model()
+    total_stored = 0
+
+    # Process in batches — encode + insert, then free vectors
+    for batch_start in range(0, len(chunks), EMBED_BATCH_SIZE):
+        batch = chunks[batch_start : batch_start + EMBED_BATCH_SIZE]
+        texts = [c["text"] for c in batch]
+
+        # Encode with adaptive batch size: group by token length to avoid
+        # OOM on long chunks (eager attention is O(batch * seq_len^2)).
+        vectors = _encode_adaptive(model, texts)
+
+        # Insert chunk metadata + vectors
+        for i, (chunk, vec) in enumerate(zip(batch, vectors)):
+            cursor = conn.execute(
+                "INSERT INTO chunks (ref_key, chunk_index, text, section_title, "
+                "granularity, start_char, text_len) VALUES (?, ?, ?, ?, ?, ?, ?)",
+                (
+                    ref_key,
+                    batch_start + i,
+                    chunk["text"],
+                    chunk["section_title"],
+                    chunk["granularity"],
+                    chunk["start_char"],
+                    len(chunk["text"]),
+                ),
+            )
+            chunk_id = cursor.lastrowid
+            conn.execute(
+                "INSERT INTO chunk_embeddings (rowid, embedding) VALUES (?, ?)",
+                (chunk_id, _serialize_f32(vec.tolist())),
+            )
+
+        conn.commit()
+        total_stored += len(batch)
+
+        # Update progress
+        conn.execute(
+            "UPDATE ref_status SET stored_chunks = ? WHERE ref_key = ?",
+            (total_stored, ref_key),
+        )
+        conn.commit()
+
+        # vectors go out of scope here — GC can reclaim
+
+    # Mark complete
+    conn.execute("UPDATE ref_status SET complete = 1 WHERE ref_key = ?", (ref_key,))
+    conn.execute(
+        "INSERT OR REPLACE INTO embed_meta (key, value) VALUES ('model', ?)",
+        (model_name,),
+    )
+    conn.commit()
+    conn.close()
+
+    return total_stored
+
+
+def retrieve_similar(
+    query_text: str,
+    ref_key: str,
+    references_dir: Path,
+    top_k: int = 15,
+) -> list[dict]:
+    """Find chunks most similar to query via sqlite-vec KNN.
+
+    Returns list of dicts with: text, section_title, granularity, distance, score.
+    Distance is cosine distance (0 = identical, 2 = opposite).
+    Score is cosine similarity (1 - distance).
+    """
+    from sciwrite_lint.references.reference_store import _get_embedding_model
+
+    db_file = _db_path(references_dir)
+    if not db_file.exists():
+        return []
+
+    conn = _open_db(references_dir)
+
+    # Check model compatibility
+    stored_model = conn.execute(
+        "SELECT value FROM embed_meta WHERE key='model'"
+    ).fetchone()
+    from sciwrite_lint.references.reference_store import _get_embedding_config
+
+    current_model, _, _ = _get_embedding_config()
+    if stored_model and stored_model[0] != current_model:
+        logger.warning(
+            "Embedding model changed ({} → {}), re-embedding needed",
+            stored_model[0],
+            current_model,
+        )
+        conn.close()
+        return []
+
+    # Encode query
+    model = _get_embedding_model()
+    query_vec = model.encode(query_text, normalize_embeddings=True)
+    query_blob = _serialize_f32(query_vec.tolist())
+
+    # KNN search: vec0 doesn't support WHERE on joined columns,
+    # so fetch top candidates globally and filter by ref_key in Python.
+    # Over-fetch to ensure enough results after filtering.
+    fetch_k = top_k * 10
+    rows = conn.execute(
+        """
+        SELECT c.text, c.section_title, c.granularity, c.start_char,
+               c.text_len, ce.distance, c.ref_key
+        FROM chunk_embeddings ce
+        INNER JOIN chunks c ON c.id = ce.rowid
+        WHERE ce.embedding MATCH ?
+            AND k = ?
+        ORDER BY ce.distance
+        """,
+        (query_blob, fetch_k),
+    ).fetchall()
+    # Filter to target reference
+    rows = [r for r in rows if r[6] == ref_key][:top_k]
+
+    conn.close()
+
+    return [
+        {
+            "text": row[0],
+            "section_title": row[1],
+            "granularity": row[2],
+            "start_char": row[3],
+            "text_len": row[4],
+            "distance": row[5],
+            "score": 1.0 - row[5],
+        }
+        for row in rows
+    ]
+
+
+def delete_embeddings(ref_key: str, references_dir: Path) -> int:
+    """Delete all embeddings for a reference. Returns count deleted."""
+    db_file = _db_path(references_dir)
+    if not db_file.exists():
+        return 0
+
+    conn = _open_db(references_dir)
+    count = _delete_ref_data(conn, ref_key)
+    conn.commit()
+    conn.close()
+    return count
+
+
+def clear_all_embeddings(references_dir: Path) -> int:
+    """Delete the entire embeddings DB. Returns number of chunks removed.
+
+    Used by ``--fresh`` to force re-embedding of all references.
+    """
+    db_file = _db_path(references_dir)
+    if not db_file.exists():
+        return 0
+
+    conn = _open_db(references_dir)
+    count = conn.execute("SELECT COUNT(*) FROM chunks").fetchone()[0]
+    conn.execute("DELETE FROM chunks")
+    conn.execute("DELETE FROM ref_status")
+    conn.execute("DELETE FROM embed_meta")
+    # Drop vec table — will be recreated with correct dim on next store
+    try:
+        conn.execute("DROP TABLE IF EXISTS chunk_embeddings")
+    except Exception as e:
+        logger.debug(f"chunk_embeddings drop skipped ({type(e).__name__}: {e})")
+    conn.commit()
+    conn.execute("VACUUM")
+    conn.close()
+    logger.info("Cleared {} embeddings from {}", count, db_file)
+    return count
+
+
+def get_stored_model(references_dir: Path) -> str | None:
+    """Return the embedding model name stored in the DB, or None."""
+    db_file = _db_path(references_dir)
+    if not db_file.exists():
+        return None
+    try:
+        conn = _open_db(references_dir)
+        row = conn.execute("SELECT value FROM embed_meta WHERE key='model'").fetchone()
+        conn.close()
+        return row[0] if row else None
+    except Exception:
+        return None