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

Integration test:
`tests/integration_tests/document_loaders/test_wikipedia.py`
pull/18603/head
Christophe Bornet 4 months ago committed by GitHub
parent 4cbfeeb1c2
commit f414f5cdb9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

@ -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)

@ -1,6 +1,6 @@
"""Util that calls Wikipedia."""
import logging
from typing import Any, Dict, List, Optional
from typing import Any, Dict, Iterator, List, Optional
from langchain_core.documents import Document
from langchain_core.pydantic_v1 import BaseModel, root_validator
@ -104,13 +104,21 @@ class WikipediaAPIWrapper(BaseModel):
Returns: a list of documents.
"""
return list(self.lazy_load(query))
def lazy_load(self, query: str) -> Iterator[Document]:
"""
Run Wikipedia search and get the article text plus the meta information.
See
Returns: a list of documents.
"""
page_titles = self.wiki_client.search(
query[:WIKIPEDIA_MAX_QUERY_LENGTH], results=self.top_k_results
)
docs = []
for page_title in page_titles[: self.top_k_results]:
if wiki_page := self._fetch_page(page_title):
if doc := self._page_to_document(page_title, wiki_page):
docs.append(doc)
return docs
yield doc

Loading…
Cancel
Save