citegraph 0.2.0__py3-none-any.whl

citegraph/extract_metadata.py

@@ -0,0 +1,81 @@
+"""Extract bibliographic metadata of source papers using Gemini."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+from citegraph.ids import make_work_id
+from citegraph.llm import GeminiClient, LLMError, parse_structured_response
+from citegraph.schemas import PaperMetadata
+
+logger = logging.getLogger(__name__)
+
+
+SYSTEM_INSTRUCTION = (
+    "You are a helpful assistant with great expertise in extracting "
+    "bibliographic information from academic papers in markdown format."
+)
+
+# Title/authors/journal/year live on the first page of essentially every
+# academic paper, so we send only a prefix to keep input tokens small.
+DEFAULT_METADATA_INPUT_CHARS = 12_000
+
+
+def _build_prompt(paper_content: str) -> str:
+    return f"""
+Extract the basic bibliographic information from the following paper. The information should include:
+- Title: The full title of the paper
+- Authors_List: A list of all authors
+- Journal: The journal or venue where the paper was published
+- Year: The publication year as an integer
+
+Make sure to not include any dirty characters that might have come from problems in the PDF.
+
+Here is the paper content:
+{paper_content}
+"""
+
+
+def extract_metadata_from_markdown(
+    markdown_path: Path | str,
+    *,
+    client: GeminiClient,
+    max_input_chars: int = DEFAULT_METADATA_INPUT_CHARS,
+) -> PaperMetadata:
+    """Extract :class:`PaperMetadata` from a single markdown-rendered paper.
+
+    Only the first ``max_input_chars`` characters of the markdown are sent to
+    the LLM. Pass ``0`` to disable truncation.
+    """
+    markdown_path = Path(markdown_path)
+    content = markdown_path.read_text(encoding="utf-8")
+    if max_input_chars and len(content) > max_input_chars:
+        logger.debug(
+            "Truncating %s from %d to %d chars for metadata extraction",
+            markdown_path.name,
+            len(content),
+            max_input_chars,
+        )
+        content = content[:max_input_chars]
+
+    response = client.generate_structured(
+        prompt=_build_prompt(content),
+        response_schema=PaperMetadata,
+        system_instruction=SYSTEM_INSTRUCTION,
+        max_output_tokens=500,
+    )
+    parsed = parse_structured_response(response, schema=PaperMetadata)
+    if not isinstance(parsed, PaperMetadata):
+        raise LLMError(
+            f"Expected PaperMetadata for {markdown_path.name}, got {type(parsed)!r}"
+        )
+    return parsed
+
+
+def metadata_to_record(meta: PaperMetadata, source_file: str) -> dict:
+    """Convert a :class:`PaperMetadata` into a flat dict ready for a DataFrame."""
+    record = meta.model_dump()
+    record["source_file"] = source_file
+    record["id"] = make_work_id(meta.Authors_List or meta.Authors, meta.Year, meta.Title)
+    return record

citegraph/extract_references.py

@@ -0,0 +1,376 @@
+"""Extract the reference list of a paper using Gemini structured output."""
+
+from __future__ import annotations
+
+import logging
+import re
+from pathlib import Path
+
+from citegraph.llm import GeminiClient, parse_structured_response, response_was_truncated
+from citegraph.schemas import Reference
+
+# Long bibliographies sometimes hit ``max_output_tokens``. When that happens we
+# retry once with a higher cap before falling back to the JSON-repair path in
+# :mod:`citegraph.llm`. The hard cap stops us from exceeding model output limits
+# on a runaway retry.
+_TRUNCATION_RETRY_MULTIPLIER = 2
+_TRUNCATION_HARD_CAP_TOKENS = 80_000
+
+# A sanity-check ratio: warn if the LLM returned fewer than this fraction of the
+# refs implied by the body's inline-citation count. Set conservatively so we
+# only flag clear shortfalls, not minor undercounts (citation-counting from
+# free text is necessarily approximate).
+_INLINE_CITATION_RATIO_THRESHOLD = 0.5
+# Skip the sanity check on bodies with very few citations — the estimate is
+# noisy at small N and the absolute gap isn't meaningful.
+_INLINE_CITATION_MIN_ESTIMATE = 10
+
+logger = logging.getLogger(__name__)
+
+
+# Strict regex for the heading of a paper's bibliography section. Matches
+# markdown headers like "## References", "# Bibliography", "### 5. References",
+# "## References and Notes", "## Works Cited". Deliberately rejects headers
+# with substantive trailing text (e.g. "## References to Prior Work") to avoid
+# slicing at the wrong place.
+_REF_KEYWORD = (
+    r"(?:references?|bibliography|works\s+cited|literature\s+cited|cited\s+literature)"
+)
+_REF_HEADER_RE = re.compile(
+    rf"^\#{{1,6}}\s+(?:[\dIVX]+\.?\s+)?{_REF_KEYWORD}(?:\s+(?:and\s+notes|cited))?\s*$",
+    re.IGNORECASE | re.MULTILINE,
+)
+
+# Any markdown header, capturing its level — used to find where a
+# references section ends.
+_ANY_HEADER_RE = re.compile(r"^(\#{1,6})\s", re.MULTILINE)
+
+# Back matter that follows a bibliography and must not be read as part of it.
+# Deliberately a *stop list* rather than "any header ends the section": the
+# splitter has to fail closed, because headers that interrupt a real reference
+# list are common — running page heads ("## 94 ECONOMIA, Spring 2009",
+# "## A. Pfaff et al.") and repeated article titles both appear mid-list in
+# scanned sources, and truncating there silently drops references. An
+# unrecognised header therefore keeps the historical run-to-end-of-file
+# behaviour. The keyword must *lead* the title (after optional numbering), so
+# "## Appendix A: Instructions" matches while "## References to Prior Work"
+# does not.
+_BACKMATTER_KEYWORD = (
+    r"(?:appendi(?:x|ces)"
+    r"|supp(?:lementary|orting)\s+(?:online\s+)?(?:material|information|data|methods?)s?"
+    r"|acknowledge?ments?"
+    r"|author\s+contributions?"
+    r"|authors?\s+(?:and|&)\s+a\S*ffiliations?"
+    r"|author\s+information"
+    r"|about\s+the\s+authors?"
+    r"|(?:end|foot)notes?"
+    r"|notes?"
+    r"|biograph(?:y|ical\s+\S+)"
+    r"|curriculum\s+vitae"
+    r"|(?:figures?|tables?)\s+(?:and|&)\s+(?:tables?|figures?)"
+    r"|index"
+    r"|notice\s+of\s+correction"
+    r"|(?:previous\s+)?(?:doctoral\s+)?the(?:sis|ses)\s+(?:in|published|publicerade)"
+    r"|previous\s+doctoral\s+theses"
+    r"|publications?"
+    r"|awards?"
+    r"|dissemination"
+    r"|data\s+availability"
+    r"|funding"
+    r"|(?:conflicts?|competing)\s+(?:of\s+)?interests?"
+    r"|declarations?)"
+)
+_BACKMATTER_HEADER_RE = re.compile(
+    rf"^\#{{1,6}}\s+(?:[\dIVXA-Z]+[.):]?\s+)?{_BACKMATTER_KEYWORD}\b",
+    re.IGNORECASE,
+)
+
+
+def _is_backmatter_header(content: str, start: int) -> bool:
+    """Whether the markdown header beginning at ``start`` opens back matter."""
+    line_end = content.find("\n", start)
+    line = content[start:] if line_end == -1 else content[start:line_end]
+    return bool(_BACKMATTER_HEADER_RE.match(line))
+
+
+def split_reference_sections(content: str) -> tuple[str, list[str]]:
+    """Return ``(body, sections)`` — all bibliography sections in ``content``.
+
+    Multi-chapter documents (PhD theses, edited volumes) carry one references
+    section per chapter; every :data:`_REF_HEADER_RE` match starts a section.
+    A non-final section ends at the next header of the same or higher level
+    (sub-headers like ``### Books`` inside a bibliography don't terminate it).
+    The final section runs to end-of-file *unless* a same-or-higher-level
+    header opens recognised back matter (:data:`_BACKMATTER_HEADER_RE`) —
+    appendices, supporting information, a department's list of previous
+    doctoral theses — which the extractor would otherwise read as references.
+    Anything unrecognised keeps the run-to-end-of-file behaviour, so a running
+    page head interrupting a reference list never truncates it. ``body`` is
+    everything outside the sections. No match returns ``(content, [])``.
+    """
+    matches = list(_REF_HEADER_RE.finditer(content))
+    if not matches:
+        return content, []
+
+    sections: list[str] = []
+    body_parts: list[str] = []
+    cursor = 0
+    for i, match in enumerate(matches):
+        start = match.start()
+        if start < cursor:
+            continue  # nested inside the previous section; already captured
+        body_parts.append(content[cursor:start])
+        end = len(content)
+        level = len(re.match(r"\#+", match.group(0)).group(0))
+        is_final = i == len(matches) - 1
+        for header in _ANY_HEADER_RE.finditer(content, match.end()):
+            if len(header.group(1)) > level:
+                continue
+            # A non-final section ends at the next section of any kind; the
+            # final one only at back matter, so unknown headers fail closed.
+            if is_final and not _is_backmatter_header(content, header.start()):
+                continue
+            end = header.start()
+            break
+        sections.append(content[start:end])
+        cursor = end
+    body_parts.append(content[cursor:])
+    return "".join(body_parts), sections
+
+
+def split_at_references_header(content: str) -> tuple[str, str | None]:
+    """Return ``(body, refs)`` with all bibliography sections in ``refs``.
+
+    Joins every section found by :func:`split_reference_sections` (a single
+    section comes back byte-identical to the raw slice). If none is found,
+    returns ``(content, None)`` so callers can decide how to react to the
+    missing header.
+    """
+    body, sections = split_reference_sections(content)
+    if not sections:
+        return content, None
+    return body, "\n\n".join(sections)
+
+
+def slice_to_references_section(content: str) -> tuple[str, bool]:
+    """Return ``(text, found)`` — content from the bibliography header onward.
+
+    If no header matches, returns the original ``content`` unchanged with
+    ``found=False`` so callers can fall back to sending the full markdown
+    rather than dropping the paper.
+    """
+    body, refs = split_at_references_header(content)
+    if refs is None:
+        return content, False
+    return refs, True
+
+
+# Numeric in-text citations: \[12\], \[12, 15\], \[12-15\], \[12, 15, 18\], etc.
+_INLINE_NUMERIC_RE = re.compile(r"\[([\d,\s\-–]+)\]")
+_NUMERIC_SPLIT_RE = re.compile(r"[,\s\-–]+")
+
+# Author-year citations in either parenthetical or narrative form:
+#   "(Smith, 2019)", "(Smith and Jones, 2019)", "(Smith et al., 2019)"
+#   "Smith (2019)", "Smith and Jones (2019)", "Smith et al. (2019)"
+# The author group is captured loosely; we lower-case it for deduplication.
+_AUTHOR = r"[A-Z][A-Za-z\-']+(?:\s+(?:and|&)\s+[A-Z][A-Za-z\-']+)?(?:\s+et\s+al\.?)?"
+_INLINE_PAREN_RE = re.compile(rf"\(({_AUTHOR}),?\s+(\d{{4}})[a-z]?\)")
+_INLINE_NARRATIVE_RE = re.compile(rf"\b({_AUTHOR})\s+\((\d{{4}})[a-z]?\)")
+
+
+def estimate_inline_citations(body: str) -> int:
+    """Approximate the number of distinct references cited inline in ``body``.
+
+    Counts:
+
+    - Unique integers inside ``[...]`` brackets (IEEE/ACM-style numeric).
+    - Unique ``(author_lowercased, year)`` tokens from author-year citations,
+      both ``(Smith, 2019)`` and narrative ``Smith (2019)`` forms.
+
+    Returns the *sum* of the two — papers typically use one style, so summing
+    gives a sensible total without much double-counting. The result is a
+    deliberately rough estimate intended for order-of-magnitude sanity checks
+    (e.g. "we extracted 12 refs but the body cites 80 — likely a parsing
+    failure"), not a precise citation count.
+    """
+    numeric: set[int] = set()
+    for match in _INLINE_NUMERIC_RE.finditer(body):
+        for token in _NUMERIC_SPLIT_RE.split(match.group(1)):
+            if token.isdigit():
+                numeric.add(int(token))
+
+    author_year: set[tuple[str, str]] = set()
+    for pattern in (_INLINE_PAREN_RE, _INLINE_NARRATIVE_RE):
+        for match in pattern.finditer(body):
+            author = re.sub(r"\s+", " ", match.group(1).lower().strip())
+            author_year.add((author, match.group(2)))
+
+    return len(numeric) + len(author_year)
+
+
+SYSTEM_INSTRUCTION = (
+    "You are an expert assistant specializing in extracting bibliographic "
+    "information from the references or bibliography section of an academic "
+    "paper. You can interpret APA, MLA, Chicago and IEEE citation formats, "
+    "and you handle complex layouts including double-column pages. Your "
+    "goal is to extract every reference comprehensively and return them in "
+    "a structured form."
+)
+
+
+def _build_prompt(paper_content: str) -> str:
+    return f"""
+# Instructions
+Please extract the complete bibliographic information for all the papers cited in the references or bibliography section of the document below. For each reference, return:
+
+- Title: The full title of the paper
+- Authors_List: A list of all authors
+- Journal: The journal or venue where the paper was published
+- Year: The publication year as an integer
+
+Take into account the following:
+
+1. Capture every reference. Do not skip any.
+2. Focus on the references / bibliography section at the end of the paper.
+3. Strip any garbled characters that may have come from PDF extraction artefacts.
+4. If a field is genuinely missing in the original, use an empty string for text fields and 0 for Year.
+
+# Paper
+{paper_content}
+"""
+
+
+def extract_references_from_markdown(
+    markdown_path: Path | str,
+    *,
+    client: GeminiClient,
+    max_output_tokens: int | None = None,
+    slice_to_references: bool = True,
+) -> list[Reference]:
+    """Extract a list of :class:`Reference` from a single markdown-rendered paper.
+
+    By default the markdown is sliced to its bibliography section before being
+    sent to the LLM (see :func:`slice_to_references_section`). Pass
+    ``slice_to_references=False`` to send the full file.
+
+    If the LLM response is cut off by ``max_output_tokens`` (detected via
+    :func:`response_was_truncated`), the call is retried once with the cap
+    raised by ``_TRUNCATION_RETRY_MULTIPLIER``, up to
+    ``_TRUNCATION_HARD_CAP_TOKENS``. A persistent truncation is logged at
+    error level and the partial result is still returned.
+    """
+    markdown_path = Path(markdown_path)
+    content = markdown_path.read_text(encoding="utf-8")
+
+    body_for_sanity_check: str | None = None
+    sections = [content]
+    if slice_to_references:
+        body, found_sections = split_reference_sections(content)
+        if found_sections:
+            logger.debug(
+                "Sliced %s to %d references section(s): %d -> %d chars",
+                markdown_path.name,
+                len(found_sections),
+                len(content),
+                sum(len(s) for s in found_sections),
+            )
+            if len(found_sections) > 1:
+                logger.info(
+                    "Found %d references sections in %s (multi-chapter document); "
+                    "extracting each separately",
+                    len(found_sections),
+                    markdown_path.name,
+                )
+            sections = found_sections
+            body_for_sanity_check = body
+        else:
+            logger.info(
+                "No references-section header found in %s; sending full markdown",
+                markdown_path.name,
+            )
+
+    initial_cap = max_output_tokens or client.default_max_output_tokens
+    parsed: list[Reference] = []
+    for idx, section in enumerate(sections):
+        label = markdown_path.name
+        if len(sections) > 1:
+            label = f"{markdown_path.name} (section {idx + 1}/{len(sections)})"
+        parsed.extend(
+            _extract_section(client=client, section=section, label=label, initial_cap=initial_cap)
+        )
+
+    if body_for_sanity_check is not None:
+        _check_extracted_count(markdown_path, body_for_sanity_check, len(parsed))
+
+    return parsed
+
+
+def _extract_section(
+    *, client: GeminiClient, section: str, label: str, initial_cap: int
+) -> list[Reference]:
+    """Run the extraction call for one references section, with truncation retry."""
+    prompt = _build_prompt(section)
+
+    def _call(cap: int):
+        return client.generate_structured(
+            prompt=prompt,
+            response_schema=list[Reference],
+            system_instruction=SYSTEM_INSTRUCTION,
+            max_output_tokens=cap,
+        )
+
+    response = _call(initial_cap)
+
+    if response_was_truncated(response):
+        retry_cap = min(
+            initial_cap * _TRUNCATION_RETRY_MULTIPLIER, _TRUNCATION_HARD_CAP_TOKENS
+        )
+        if retry_cap > initial_cap:
+            logger.warning(
+                "References response for %s was truncated at %d tokens; "
+                "retrying with cap %d",
+                label,
+                initial_cap,
+                retry_cap,
+            )
+            response = _call(retry_cap)
+            if response_was_truncated(response):
+                logger.error(
+                    "References response for %s still truncated at %d tokens; "
+                    "extracted list will be incomplete",
+                    label,
+                    retry_cap,
+                )
+        else:
+            logger.error(
+                "References response for %s truncated at hard cap %d tokens; "
+                "cannot retry — extracted list will be incomplete",
+                label,
+                initial_cap,
+            )
+
+    parsed = parse_structured_response(response, schema=Reference)
+    if not isinstance(parsed, list):
+        logger.warning("Expected a list of references for %s, got %r", label, type(parsed))
+        return []
+    return parsed
+
+
+def _check_extracted_count(
+    markdown_path: Path, body: str, extracted_count: int
+) -> None:
+    """Warn if the extracted count is far below the body's inline-citation estimate."""
+    estimate = estimate_inline_citations(body)
+    if estimate < _INLINE_CITATION_MIN_ESTIMATE:
+        return
+    if extracted_count >= estimate * _INLINE_CITATION_RATIO_THRESHOLD:
+        return
+    logger.warning(
+        "Possible references undercount for %s: extracted %d but body cites ~%d "
+        "distinct works (threshold %.0f%%). The bibliography may have been "
+        "partially parsed or the LLM may have skipped entries.",
+        markdown_path.name,
+        extracted_count,
+        estimate,
+        _INLINE_CITATION_RATIO_THRESHOLD * 100,
+    )