mirror of
https://github.com/Mintplex-Labs/langchain-python.git
synced 2026-08-26 16:55:38 -04:00
Eugene Yurtsev
2ceb807da2
# Add PDF parser implementations
This PR separates the data loading from the parsing for a number of
existing PDF loaders.
Parser tests have been designed to help encourage developers to create a
consistent interface for parsing PDFs.
This interface can be made more consistent in the future by adding
information into the initializer on desired behavior with respect to splitting by
page etc.
This code is expected to be backwards compatible -- with the exception
of a bug fix with pymupdf parser which was returning `bytes` in the page
content rather than strings.
Also changing the lazy parser method of document loader to return an
Iterator rather than Iterable over documents.
## Before submitting
<!-- If you're adding a new integration, include an integration test and
an example notebook showing its use! -->
## Who can review?
Community members can review the PR once tests pass. Tag
maintainers/contributors who might be interested:
@
<!-- For a quicker response, figure out the right person to tag with @
@hwchase17 - project lead
Tracing / Callbacks
- @agola11
Async
- @agola11
DataLoader Abstractions
- @eyurtsev
LLM/Chat Wrappers
- @hwchase17
- @agola11
Tools / Toolkits
- @vowelparrot
-->
48 lines
1.6 KiB
Python
48 lines
1.6 KiB
Python
import json
|
|
from pathlib import Path
|
|
from typing import Iterator, List, Union
|
|
|
|
from langchain.docstore.document import Document
|
|
from langchain.document_loaders.base import BaseLoader
|
|
|
|
|
|
class TomlLoader(BaseLoader):
|
|
"""
|
|
A TOML document loader that inherits from the BaseLoader class.
|
|
|
|
This class can be initialized with either a single source file or a source
|
|
directory containing TOML files.
|
|
"""
|
|
|
|
def __init__(self, source: Union[str, Path]):
|
|
"""Initialize the TomlLoader with a source file or directory."""
|
|
self.source = Path(source)
|
|
|
|
def load(self) -> List[Document]:
|
|
"""Load and return all documents."""
|
|
return list(self.lazy_load())
|
|
|
|
def lazy_load(self) -> Iterator[Document]:
|
|
"""Lazily load the TOML documents from the source file or directory."""
|
|
import tomli
|
|
|
|
if self.source.is_file() and self.source.suffix == ".toml":
|
|
files = [self.source]
|
|
elif self.source.is_dir():
|
|
files = list(self.source.glob("**/*.toml"))
|
|
else:
|
|
raise ValueError("Invalid source path or file type")
|
|
|
|
for file_path in files:
|
|
with file_path.open("r", encoding="utf-8") as file:
|
|
content = file.read()
|
|
try:
|
|
data = tomli.loads(content)
|
|
doc = Document(
|
|
page_content=json.dumps(data),
|
|
metadata={"source": str(file_path)},
|
|
)
|
|
yield doc
|
|
except tomli.TOMLDecodeError as e:
|
|
print(f"Error parsing TOML file {file_path}: {e}")
|