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