mirror of
https://github.com/hwchase17/langchain.git
synced 2025-08-12 06:13:36 +00:00
community[minor]: Implement lazy_load() for WikipediaLoader (#18680)
Integration test: `tests/integration_tests/document_loaders/test_wikipedia.py`
This commit is contained in:
parent
4cbfeeb1c2
commit
f414f5cdb9
@ -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…
Reference in New Issue
Block a user