community[minor]: Implement lazy_load() for WikipediaLoader (#18680)

Integration test:
`tests/integration_tests/document_loaders/test_wikipedia.py`
This commit is contained in:
Christophe Bornet
2024-03-06 19:03:21 +01:00
committed by GitHub
parent 4cbfeeb1c2
commit f414f5cdb9
2 changed files with 16 additions and 9 deletions

View File

@@ -1,4 +1,4 @@
from typing import List, Optional
from typing import Iterator, Optional
from langchain_core.documents import Document
@@ -42,12 +42,12 @@ class WikipediaLoader(BaseLoader):
self.load_all_available_meta = load_all_available_meta
self.doc_content_chars_max = doc_content_chars_max
def load(self) -> List[Document]:
def lazy_load(self) -> Iterator[Document]:
"""
Loads the query result from Wikipedia into a list of Documents.
Returns:
List[Document]: A list of Document objects representing the loaded
A list of Document objects representing the loaded
Wikipedia pages.
"""
client = WikipediaAPIWrapper(
@@ -56,5 +56,4 @@ class WikipediaLoader(BaseLoader):
load_all_available_meta=self.load_all_available_meta,
doc_content_chars_max=self.doc_content_chars_max,
)
docs = client.load(self.query)
return docs
yield from client.load(self.query)