From 8e10ac422e4e6b193fc35e1d64d7f0c5208faa8d Mon Sep 17 00:00:00 2001 From: Matt Robinson Date: Thu, 27 Apr 2023 01:09:45 -0400 Subject: [PATCH] enhancement: add elements mode to `UnstructuredURLLoader` (#3456) ### Summary Updates the `UnstructuredURLLoader` to include a "elements" mode that retains additional metadata from `unstructured`. This makes `UnstructuredURLLoader` consistent with other unstructured loaders, which also support "elements" mode. Patched mode into the existing `UnstructuredURLLoader` class instead of inheriting from `UnstructuredBaseLoader` because it significantly simplified the implementation. ### Testing This should still work and show the url in the source for the metadata ```python from langchain.document_loaders import UnstructuredURLLoader urls = ["https://www.understandingwar.org/sites/default/files/Russian%20Offensive%20Campaign%20Assessment%2C%20April%2011%2C%202023.pdf"] loader = UnstructuredURLLoader(urls=urls, headers={"Accept": "application/json"}, strategy="fast") docs = loader.load() print(docs[0].page_content[:1000]) docs[0].metadata ``` This should now work and show additional metadata from `unstructured`. This should still work and show the url in the source for the metadata ```python from langchain.document_loaders import UnstructuredURLLoader urls = ["https://www.understandingwar.org/sites/default/files/Russian%20Offensive%20Campaign%20Assessment%2C%20April%2011%2C%202023.pdf"] loader = UnstructuredURLLoader(urls=urls, headers={"Accept": "application/json"}, strategy="fast", mode="elements") docs = loader.load() print(docs[0].page_content[:1000]) docs[0].metadata ``` --- langchain/document_loaders/url.py | 25 ++++++++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/langchain/document_loaders/url.py b/langchain/document_loaders/url.py index 3f52c2b8..631052f6 100644 --- a/langchain/document_loaders/url.py +++ b/langchain/document_loaders/url.py @@ -15,6 +15,7 @@ class UnstructuredURLLoader(BaseLoader): self, urls: List[str], continue_on_failure: bool = True, + mode: str = "single", **unstructured_kwargs: Any, ): """Initialize with file path.""" @@ -29,6 +30,9 @@ class UnstructuredURLLoader(BaseLoader): "`pip install unstructured`" ) + self._validate_mode(mode) + self.mode = mode + headers = unstructured_kwargs.pop("headers", {}) if len(headers.keys()) != 0: warn_about_headers = False @@ -48,6 +52,13 @@ class UnstructuredURLLoader(BaseLoader): self.headers = headers self.unstructured_kwargs = unstructured_kwargs + def _validate_mode(self, mode: str) -> None: + _valid_modes = {"single", "elements"} + if mode not in _valid_modes: + raise ValueError( + f"Got {mode} for `mode`, but should be one of `{_valid_modes}`" + ) + def __is_headers_available_for_html(self) -> bool: _unstructured_version = self.__version.split("-")[0] unstructured_version = tuple([int(x) for x in _unstructured_version.split(".")]) @@ -94,7 +105,15 @@ class UnstructuredURLLoader(BaseLoader): continue else: raise e - text = "\n\n".join([str(el) for el in elements]) - metadata = {"source": url} - docs.append(Document(page_content=text, metadata=metadata)) + + if self.mode == "single": + text = "\n\n".join([str(el) for el in elements]) + metadata = {"source": url} + docs.append(Document(page_content=text, metadata=metadata)) + elif self.mode == "elements": + for element in elements: + metadata = element.metadata.to_dict() + metadata["category"] = element.category + docs.append(Document(page_content=str(element), metadata=metadata)) + return docs