citegraph 0.2.0__py3-none-any.whl

citegraph/jel.py

@@ -0,0 +1,167 @@
+"""JEL classification codes of the core papers.
+
+Most economics papers print their JEL codes on the title page, and docling
+carries that line into the markdown, so the codes are read back with a
+regex rather than another LLM call: adding a field to ``PaperMetadata``
+would invalidate the metadata cache, and re-extracting every paper could
+shift titles, hence work ids, hence every id-keyed curated file. Papers
+whose PDF carries no codes (most working papers and chapters) get them by hand through
+the ``JEL_Codes`` column of ``work_annotations.csv``; a hand-assigned value
+overrides extraction, because a reviewer who opened the PDF beats a regex.
+
+``core_jel_codes`` is the merged view with provenance, and ``jel_codes.csv``
+is it on disk. The general-category names come from the AEA's
+``classifications.xml`` (the ``X00`` entries, minus their ``: General`` tail).
+"""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+import pandas as pd
+
+from citegraph.annotations import load_annotations
+from citegraph.io import OutLayout
+
+JEL_GENERAL_CATEGORIES: dict[str, str] = {
+    "A": "General Economics and Teaching",
+    "B": "History of Economic Thought, Methodology, and Heterodox Approaches",
+    "C": "Mathematical and Quantitative Methods",
+    "D": "Microeconomics",
+    "E": "Macroeconomics and Monetary Economics",
+    "F": "International Economics",
+    "G": "Financial Economics",
+    "H": "Public Economics",
+    "I": "Health, Education, and Welfare",
+    "J": "Labor and Demographic Economics",
+    "K": "Law and Economics",
+    "L": "Industrial Organization",
+    "M": "Business Administration and Business Economics; Marketing; Accounting; "
+    "Personnel Economics",
+    "N": "Economic History",
+    "O": "Economic Development, Innovation, Technological Change, and Growth",
+    "P": "Political Economy and Comparative Economic Systems",
+    "Q": "Agricultural and Natural Resource Economics; Environmental and Ecological Economics",
+    "R": "Urban, Rural, Regional, Real Estate, and Transportation Economics",
+    "Y": "Miscellaneous Categories",
+    "Z": "Other Special Topics",
+}
+
+# docling leaves font escapes in the text; these are the ones that land in JEL lines.
+_ARTEFACTS = {"/uni00A0": " ", "/C1": " ", "/uniFB01": "fi", "/fication": "fication"}
+
+# "JEL", an optional label ("classification", "codes", "classification numbers",
+# with the fi-ligature mangled every way docling manages), then the rest of the line.
+_LABEL_RE = re.compile(
+    r"\bJEL\b[ \t]*[-–:]?[ \t]*"
+    r"(?:classi\s*fi?\s*c?ations?|codes?|numbers?)?"
+    r"(?:[ \t]+(?:numbers?|codes?))?"
+    r"[ \t]*[:.]?[ \t]*(?P<rest>[^\n]*)",
+    re.IGNORECASE,
+)
+_CODE_RE = re.compile(r"\b([A-Z0]) ?(\d{1,2})\b")
+_FILLER_RE = re.compile(r"\band\b|[\s,;.·/()\[\]\-–]")
+_CODE_FORMAT_RE = re.compile(r"^[A-Z]\d{1,2}$")
+
+
+def _codes_in(chunk: str) -> list[str]:
+    """Codes in ``chunk`` when it is a code list and nothing else."""
+    codes = [
+        ("O" if letter == "0" else letter) + digits for letter, digits in _CODE_RE.findall(chunk)
+    ]
+    codes = [c for c in codes if c[0] in JEL_GENERAL_CATEGORIES]
+    if not codes or _FILLER_RE.sub("", _CODE_RE.sub("", chunk)).strip():
+        return []
+    return codes
+
+
+def extract_jel_codes(markdown: str) -> list[str]:
+    """JEL codes printed in ``markdown``, in order, deduplicated.
+
+    Every ``JEL`` block is read — a thesis prints one per chapter — and a
+    block's codes are either on the label's own line or in the paragraphs
+    right after it, read until the first paragraph that is not a code list
+    (an Elsevier "Article info" box prints one code per paragraph). A line
+    that mentions JEL and then says anything else is prose, not a code list,
+    and contributes nothing.
+    """
+    text = markdown
+    for artefact, replacement in _ARTEFACTS.items():
+        text = text.replace(artefact, replacement)
+    found: dict[str, None] = {}
+    for match in _LABEL_RE.finditer(text):
+        rest = match.group("rest").strip()
+        if rest:
+            found.update(dict.fromkeys(_codes_in(rest)))
+            continue
+        for paragraph in re.split(r"\n\s*\n", text[match.end() :].lstrip("\n")):
+            codes = _codes_in(paragraph)
+            if not codes:
+                break
+            found.update(dict.fromkeys(codes))
+    return list(found)
+
+
+def _parse_manual(value: object, work_id: str) -> list[str]:
+    if value is None or pd.isna(value) or not str(value).strip():
+        return []
+    tokens = [token for token in re.split(r"[;,\s]+", str(value).strip()) if token]
+    for token in tokens:
+        code = token.upper()
+        if not _CODE_FORMAT_RE.match(code) or code[0] not in JEL_GENERAL_CATEGORIES:
+            raise ValueError(
+                f"work_annotations.csv: work {work_id!r} has JEL_Codes entry {token!r}; "
+                "expected a letter and one or two digits, e.g. D71 or Q2"
+            )
+    return list(dict.fromkeys(token.upper() for token in tokens))
+
+
+def core_jel_codes(out_dir: str | Path) -> pd.DataFrame:
+    """One row per core work: ``JEL_Codes``, ``JEL_General`` and where they came from.
+
+    ``JEL_Source`` is ``manual`` (from ``work_annotations.csv``), ``extracted``
+    (from the markdown) or ``missing``. ``JEL_General`` is the sorted set of
+    first letters — the level the composition figure plots.
+    """
+    layout = OutLayout(out_dir)
+    works = pd.read_csv(layout.works_csv, index_col="id", keep_default_na=False)
+    core = works[works["ring"] == 0]
+    annotations = load_annotations(
+        layout.work_annotations_csv, works=works, schema_path=layout.annotation_schema_csv
+    )
+    manual = (
+        annotations["JEL_Codes"] if "JEL_Codes" in annotations.columns else pd.Series(dtype=object)
+    )
+
+    rows = []
+    for work_id, work in core.iterrows():
+        codes = _parse_manual(manual.get(work_id), str(work_id))
+        source = "manual"
+        if not codes:
+            markdown = layout.markdown_dir / str(work["source_file"])
+            if work["source_file"] and markdown.exists():
+                codes = extract_jel_codes(markdown.read_text(encoding="utf-8", errors="replace"))
+            source = "extracted" if codes else "missing"
+        rows.append(
+            {
+                "id": work_id,
+                "source_file": work["source_file"],
+                "Title": work["Title"],
+                "JEL_Codes": codes,
+                "JEL_General": sorted({code[0] for code in codes}),
+                "JEL_Source": source,
+            }
+        )
+    columns = ["id", "source_file", "Title", "JEL_Codes", "JEL_General", "JEL_Source"]
+    return pd.DataFrame(rows, columns=columns).set_index("id")
+
+
+def write_jel_codes(out_dir: str | Path) -> Path:
+    """Write ``jel_codes.csv`` (lists joined by ``;``) and return its path."""
+    frame = core_jel_codes(out_dir)
+    for column in ("JEL_Codes", "JEL_General"):
+        frame[column] = frame[column].map(";".join)
+    path = OutLayout(out_dir).jel_codes_csv
+    frame.to_csv(path)
+    return path

citegraph/journals.py

@@ -0,0 +1,182 @@
+"""Canonical journal names.
+
+The same journal reaches us under many spellings — the AER under six in one
+corpus, *Ecological Economics* under three — so any rollup that groups on the
+raw string understates concentration. Enrichment already collapses most of the
+variation by preferring the provider's container title; this module folds what
+is left.
+
+Two deliberate limits:
+
+- Folding is conservative. It normalises case, diacritics, punctuation,
+  ``&``/``and`` and a leading ``The``, and nothing else. Abbreviations are
+  *not* matched automatically: an initials heuristic over one corpus paired
+  "public choice" with "protecting the commons" and "energy policy" with
+  "economics and philosophy". Abbreviations belong in the alias file, where a
+  human vouches for them; there are only about ten worth writing.
+- Preprint repositories are not journals. OpenAlex reports RePEc, SSRN and
+  publisher eBook platforms as container titles, which would rank RePEc third
+  and SSRN ninth among this corpus's "journals". They fold to
+  :data:`WORKING_PAPER` instead, and an alias can override that.
+"""
+
+from __future__ import annotations
+
+import csv
+import math
+import re
+import unicodedata
+from collections import Counter
+from collections.abc import Iterable, Mapping, Sequence
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+WORKING_PAPER = "Working paper"
+
+# Matched against the folded key, so punctuation and case are already gone.
+_REPOSITORY_PATTERNS = (
+    "repec",
+    "research papers in economics",
+    "ssrn",
+    "ebooks",
+    "social science research network",
+)
+
+_LEADING_ARTICLE = re.compile(r"^the\s+")
+_AMPERSAND = re.compile(r"\s*&\s*")
+_NON_ALNUM = re.compile(r"[^a-z0-9]+")
+
+
+def journal_key(name: object) -> str:
+    """Fold one journal name to its grouping key.
+
+    Returns ``""`` for anything missing, so a work with no journal never
+    joins a group.
+    """
+    if name is None or (isinstance(name, float) and math.isnan(name)):
+        return ""
+    text = str(name).strip()
+    if not text:
+        return ""
+    text = unicodedata.normalize("NFKD", text)
+    text = "".join(c for c in text if not unicodedata.combining(c))
+    text = _AMPERSAND.sub(" and ", text.lower())
+    text = _LEADING_ARTICLE.sub("", text)
+    return _NON_ALNUM.sub(" ", text).strip()
+
+
+def _is_repository(key: str) -> bool:
+    return any(pattern in key for pattern in _REPOSITORY_PATTERNS)
+
+
+def canonicalize_journals(
+    names: Iterable[object],
+    aliases: Mapping[str, str] | None = None,
+    *,
+    strict: bool = False,
+) -> dict[str, str]:
+    """Map every observed journal spelling to one canonical display name.
+
+    The display name is the most frequent spelling in ``names``, with the
+    longest winning a tie — the same "fullest, most frequent wins" rule the
+    author stage uses, so a rare lowercased variant never names a group.
+
+    ``aliases`` is applied last and matched by folded key, so the user writes
+    one spelling ("Am. Econ. Rev.") and every variant of it follows. With
+    ``strict``, an alias that matches nothing raises: a curated row that
+    silently does not apply is worse than no row at all.
+    """
+    observed = [str(n).strip() for n in names if journal_key(n)]
+    counts: dict[str, Counter] = {}
+    for raw in observed:
+        counts.setdefault(journal_key(raw), Counter())[raw] += 1
+
+    alias_by_key = {journal_key(raw): canonical for raw, canonical in (aliases or {}).items()}
+    if strict:
+        unused = unused_aliases(aliases, observed)
+        if unused:
+            raise ValueError(unused_alias_message(unused))
+
+    display: dict[str, str] = {}
+    for key, spellings in counts.items():
+        if key in alias_by_key:
+            display[key] = alias_by_key[key]
+        elif _is_repository(key):
+            display[key] = WORKING_PAPER
+        else:
+            display[key] = max(spellings.items(), key=lambda kv: (kv[1], len(kv[0])))[0]
+
+    return {
+        (str(n).strip() if journal_key(n) else ""): display.get(journal_key(n), "")
+        for n in names
+    }
+
+
+def unused_aliases(
+    aliases: Mapping[str, str] | None, *vocabularies: Iterable[object]
+) -> list[str]:
+    """Curated rows whose folded key appears in none of ``vocabularies``.
+
+    A corpus spells its journals twice — once as the bibliography extracted
+    them, once as the provider reports them — and a row is only dead if it
+    matches *neither*. Ask this against every vocabulary the corpus has rather
+    than against one fold, or a row is rejected for the very thing it does:
+    an alias expanding "Ecol. Econ." matches nothing among provider titles
+    precisely because the provider already says "Ecological Economics".
+    """
+    known: set[str] = set()
+    for names in vocabularies:
+        known.update(key for key in (journal_key(n) for n in names) if key)
+    return sorted(raw for raw in (aliases or {}) if journal_key(raw) not in known)
+
+
+def unused_alias_message(unused: Sequence[str]) -> str:
+    """The one wording for a dead curated row, shared by every caller."""
+    return (
+        "Journal aliases match no journal in this corpus: "
+        + ", ".join(repr(u) for u in unused)
+        + ". Fix the spelling or delete the row."
+    )
+
+
+def load_journal_aliases(path: Path | str | None) -> dict[str, str]:
+    """Read ``journal_aliases.csv`` (``raw,canonical``); absent file is empty."""
+    if path is None:
+        return {}
+    path = Path(path)
+    if not path.exists():
+        return {}
+    with path.open(encoding="utf-8", newline="") as handle:
+        rows = list(csv.DictReader(handle))
+    aliases: dict[str, str] = {}
+    for row in rows:
+        raw = (row.get("raw") or "").strip()
+        canonical = (row.get("canonical") or "").strip()
+        if raw and canonical:
+            aliases[raw] = canonical
+    return aliases
+
+
+def add_canonical_journals(
+    df: pd.DataFrame,
+    aliases: Mapping[str, str] | None = None,
+    *,
+    strict: bool = False,
+) -> pd.DataFrame:
+    """Return ``df`` with a ``Journal_Canonical`` column beside ``Journal``.
+
+    The observed name is never overwritten — the canonical form is what you
+    group on, the observed one is what the record actually said.
+    """
+    out = df.copy()
+    if "Journal" not in out.columns:
+        out["Journal_Canonical"] = ""
+        return out
+    mapping = canonicalize_journals(out["Journal"].tolist(), aliases, strict=strict)
+    out["Journal_Canonical"] = [
+        mapping.get(str(n).strip() if journal_key(n) else "", "") for n in out["Journal"]
+    ]
+    return out

citegraph/llm.py

@@ -0,0 +1,186 @@
+"""Thin Gemini wrapper used by the extraction modules.
+
+We keep all of the provider-specific code in this single file so that the
+rest of the package only deals with Pydantic models and plain Python
+objects. The wrapper handles:
+
+- Lazy client construction with a configurable API key/model.
+- Retry/backoff via :mod:`tenacity` for transient errors.
+- A :func:`fix_incomplete_json_string` JSON repair fallback for the rare
+  cases where Gemini truncates a response.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import Any
+
+from pydantic import BaseModel
+from tenacity import (
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_exponential,
+)
+
+from citegraph.config import Settings, get_settings
+
+logger = logging.getLogger(__name__)
+
+
+def fix_incomplete_json_string(input_string: str) -> str | None:
+    """Attempt to recover a truncated JSON-array string.
+
+    The current Gemini structured-output path almost always returns valid
+    JSON, but on long bibliographies the response can hit the output-token
+    limit and arrive truncated. This function:
+
+    1. Strips surrounding single quotes (a quirk we sometimes see).
+    2. Cuts the string at the last complete ``}``.
+    3. Closes the trailing ``]`` so the result parses as a JSON array.
+
+    Returns ``None`` if no complete object boundary is found.
+    """
+    cleaned = input_string
+    if cleaned.startswith("'"):
+        cleaned = cleaned[1:]
+    if cleaned.endswith("'"):
+        cleaned = cleaned[:-1]
+
+    last_obj_end = cleaned.rfind("}")
+    if last_obj_end == -1:
+        return None
+    return cleaned[: last_obj_end + 1] + "\n]"
+
+
+class LLMError(RuntimeError):
+    """Raised when the LLM call fails after all retries."""
+
+
+def response_was_truncated(response: Any) -> bool:
+    """Return True if the Gemini response was cut off by ``max_output_tokens``.
+
+    Inspects ``response.candidates[*].finish_reason`` and matches the value
+    "MAX_TOKENS" (the SDK may expose it as an enum or a plain string). Returns
+    ``False`` defensively when the response shape is unexpected — callers can
+    fall back to other signals (e.g. JSON parse failure) in that case.
+    """
+    try:
+        candidates = getattr(response, "candidates", None) or []
+        for cand in candidates:
+            reason = getattr(cand, "finish_reason", None)
+            name = getattr(reason, "name", None) or str(reason or "")
+            if "MAX_TOKENS" in name:
+                return True
+    except Exception:  # noqa: BLE001
+        return False
+    return False
+
+
+class GeminiClient:
+    """Minimal Gemini client used by the extraction modules."""
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        model: str | None = None,
+        settings: Settings | None = None,
+    ) -> None:
+        self._settings = settings or get_settings()
+        self._api_key = api_key or self._settings.require_api_key()
+        self.model = model or self._settings.citegraph_model
+        self._client: Any | None = None
+
+    @property
+    def client(self) -> Any:
+        if self._client is None:
+            from google import genai
+
+            self._client = genai.Client(api_key=self._api_key)
+        return self._client
+
+    @property
+    def default_max_output_tokens(self) -> int:
+        """Output-token cap used when ``generate_structured`` gets ``None``."""
+        return self._settings.citegraph_max_output_tokens
+
+    @retry(
+        retry=retry_if_exception_type(Exception),
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=2, min=2, max=30),
+        reraise=True,
+    )
+    def _generate(self, *, contents: str, config: dict[str, Any]) -> Any:
+        return self.client.models.generate_content(
+            model=self.model,
+            contents=contents,
+            config=config,
+        )
+
+    def generate_structured(
+        self,
+        *,
+        prompt: str,
+        response_schema: Any,
+        system_instruction: str | None = None,
+        max_output_tokens: int | None = None,
+    ) -> Any:
+        """Run a structured-output Gemini call and return the response object.
+
+        ``response_schema`` is passed through to Gemini unchanged. Use a
+        Pydantic ``BaseModel`` subclass for a single-object response or a
+        list type (e.g. ``list[Reference]``) for an array response.
+        """
+        config: dict[str, Any] = {
+            "response_mime_type": "application/json",
+            "response_schema": response_schema,
+            "max_output_tokens": (
+                max_output_tokens or self._settings.citegraph_max_output_tokens
+            ),
+        }
+        if system_instruction:
+            config["system_instruction"] = system_instruction
+        try:
+            return self._generate(contents=prompt, config=config)
+        except Exception as exc:  # noqa: BLE001
+            raise LLMError(f"Gemini call failed: {exc}") from exc
+
+
+def parse_structured_response(
+    response: Any,
+    schema: type[BaseModel] | None = None,
+) -> Any:
+    """Parse a Gemini response, falling back to JSON repair if needed.
+
+    - If ``response.parsed`` is populated, return it.
+    - Otherwise try ``json.loads(response.text)``.
+    - On failure, attempt :func:`fix_incomplete_json_string` and retry.
+
+    When ``schema`` is provided, list elements / single objects are
+    validated through that Pydantic model before being returned.
+    """
+    parsed = getattr(response, "parsed", None)
+    if parsed is not None:
+        data = parsed
+    else:
+        raw = getattr(response, "text", None)
+        if not raw:
+            raise LLMError("Empty Gemini response (no .parsed and no .text).")
+
+        try:
+            data = json.loads(raw)
+        except json.JSONDecodeError as exc:
+            logger.warning("Gemini returned invalid JSON, attempting repair: %s", exc)
+            repaired = fix_incomplete_json_string(raw)
+            if repaired is None:
+                raise LLMError("Could not repair truncated JSON response.") from exc
+            data = json.loads(repaired)
+
+    if schema is None:
+        return data
+
+    if isinstance(data, list):
+        return [schema.model_validate(item) for item in data]
+    return schema.model_validate(data)