codekingpro/portable-devtools

1import logging2import re3from pathlib import Path4from typing import Any, Iterator, List, Mapping, Optional5 6from langchain_core.documents import Document7 8from langchain_community.document_loaders.base import BaseLoader9from langchain_community.utilities.bibtex import BibtexparserWrapper10 11logger = logging.getLogger(__name__)12 13 14class BibtexLoader(BaseLoader):15    """Load a `bibtex` file.16 17    Each document represents one entry from the bibtex file.18 19    If a PDF file is present in the `file` bibtex field, the original PDF20    is loaded into the document text. If no such file entry is present,21    the `abstract` field is used instead.22    """23 24    def __init__(25        self,26        file_path: str,27        *,28        parser: Optional[BibtexparserWrapper] = None,29        max_docs: Optional[int] = None,30        max_content_chars: Optional[int] = 4_000,31        load_extra_metadata: bool = False,32        file_pattern: str = r"[^:]+\.pdf",33    ):34        """Initialize the BibtexLoader.35 36        Args:37            file_path: Path to the bibtex file.38            parser: The parser to use. If None, a default parser is used.39            max_docs: Max number of associated documents to load. Use -1 means40                           no limit.41            max_content_chars: Maximum number of characters to load from the PDF.42            load_extra_metadata: Whether to load extra metadata from the PDF.43            file_pattern: Regex pattern to match the file name in the bibtex.44        """45        self.file_path = file_path46        self.parser = parser or BibtexparserWrapper()47        self.max_docs = max_docs48        self.max_content_chars = max_content_chars49        self.load_extra_metadata = load_extra_metadata50        self.file_regex = re.compile(file_pattern)51 52    def _load_entry(self, entry: Mapping[str, Any]) -> Optional[Document]:53        import fitz54 55        parent_dir = Path(self.file_path).parent56        # regex is useful for Zotero flavor bibtex files57        file_names = self.file_regex.findall(entry.get("file", ""))58        if not file_names:59            return None60        texts: List[str] = []61        for file_name in file_names:62            try:63                with fitz.open(parent_dir / file_name) as f:64                    texts.extend(page.get_text() for page in f)65            except FileNotFoundError as e:66                logger.debug(e)67        content = "\n".join(texts) or entry.get("abstract", "")68        if self.max_content_chars:69            content = content[: self.max_content_chars]70        metadata = self.parser.get_metadata(entry, load_extra=self.load_extra_metadata)71        return Document(72            page_content=content,73            metadata=metadata,74        )75 76    def lazy_load(self) -> Iterator[Document]:77        """Load bibtex file using bibtexparser and get the article texts plus the78        article metadata.79        See https://bibtexparser.readthedocs.io/en/master/80 81        Returns:82            a list of documents with the document.page_content in text format83        """84        try:85            import fitz  # noqa: F40186        except ImportError:87            raise ImportError(88                "PyMuPDF package not found, please install it with "89                "`pip install pymupdf`"90            )91 92        entries = self.parser.load_bibtex_entries(self.file_path)93        if self.max_docs:94            entries = entries[: self.max_docs]95        for entry in entries:96            doc = self._load_entry(entry)97            if doc:98                yield doc99