components/reference_extractor.py

"""
Reference extraction from PDF and text documents using Google Gemini AI.
"""

from google import genai
from google.genai import types
import pathlib
import json
import logging
from typing import Dict, List, Any

logger = logging.getLogger(__name__)


class ReferenceExtractor:
    """
    Extracts academic references from PDF files and text using AI.
    """

    # Shared extraction prompt template
    EXTRACTION_PROMPT_TEMPLATE = """
Identify the title, first author, year, and journal of each reference.
The references may be in various formats, including APA, MLA, and Chicago.
The references may also include DOIs, URLs, and other identifiers.

IMPORTANT: For each reference, analyze its content to determine the optimal API order:
- If the reference mentions "arXiv", "preprint", or contains "arxiv.org", put 'arxiv' first in the API order.
- If it mentions a specific journal, contains a DOI, or is clearly a published journal article, put 'crossref' first.
- For newer publications, open access articles, or general academic works, put 'openalex' first.

Format your response as JSON with these fields:
- title: The title of the academic work in the reference
- first_author: The first author's name
- year: Publication year as string
- journal: Journal or publication venue name (optional)
- reference: The full reference text without the number of the reference
- check_api_order: Array of API names to check in order ['crossref', 'arxiv', 'openalex'] based on the reference type

Output only json
"""

    def __init__(self, client: genai.Client, model: str = "gemini-2.5-flash-lite"):
        """
        Initialize the reference extractor.

        Args:
            client: Google GenAI client instance
            model: Model name to use for extraction
        """
        self.client = client
        self.model = model

    def extract_from_pdf(self, file_path: str) -> List[Dict[str, Any]]:
        """
        Extracts references from a PDF file using GenAI.

        Args:
            file_path: The path to the PDF file to extract references from.

        Returns:
            A list of extracted references.

        Raises:
            FileNotFoundError: If PDF file doesn't exist
            ValueError: If extraction fails
        """
        file_path_obj = pathlib.Path(file_path)

        if not file_path_obj.exists():
            logger.error(f"PDF file not found: {file_path}")
            raise FileNotFoundError(f"PDF file not found: {file_path}")

        prompt = f"From the PDF file, extract the references.\n{self.EXTRACTION_PROMPT_TEMPLATE}"

        try:
            logger.info(f"Extracting references from PDF: {file_path_obj.name}")

            response = self.client.models.generate_content(
                model=self.model,
                contents=[
                    types.Part.from_bytes(
                        data=file_path_obj.read_bytes(),
                        mime_type='application/pdf',
                    ),
                    prompt
                ],
                config=types.GenerateContentConfig(
                    response_mime_type="application/json",
                    temperature=0
                )
            )

            references = json.loads(response.candidates[0].content.parts[0].text)
            logger.info(f"Extracted {len(references) if isinstance(references, list) else 1} references")
            return references

        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse GenAI response as JSON: {e}")
            raise ValueError(f"Invalid JSON response from GenAI: {e}")
        except Exception as e:
            logger.exception(f"Failed to extract references from PDF: {e}")
            raise

    def extract_from_text(self, text: str) -> List[Dict[str, Any]]:
        """
        Extracts references from a text string using GenAI.

        Args:
            text: The text to extract references from.

        Returns:
            A list of extracted references.

        Raises:
            ValueError: If text is empty or extraction fails
        """
        if not text or not text.strip():
            logger.error("Empty text provided for reference extraction")
            raise ValueError("Text cannot be empty")

        prompt = f"""
From the text, extract the references.

Text:
{text}

{self.EXTRACTION_PROMPT_TEMPLATE}
"""

        try:
            logger.info(f"Extracting references from text ({len(text)} characters)")

            response = self.client.models.generate_content(
                model=self.model,
                contents=[prompt],
                config=types.GenerateContentConfig(
                    response_mime_type="application/json",
                    temperature=0
                )
            )

            references = json.loads(response.candidates[0].content.parts[0].text)
            logger.info(f"Extracted {len(references) if isinstance(references, list) else 1} references")
            return references

        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse GenAI response as JSON: {e}")
            raise ValueError(f"Invalid JSON response from GenAI: {e}")
        except Exception as e:
            logger.exception(f"Failed to extract references from text: {e}")
            raise