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,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