llama-index

1
"""Simple reader that reads wikipedia."""
2

3
from typing import Any, List
4

5
from llama_index.legacy.readers.base import BasePydanticReader
6
from llama_index.legacy.schema import Document
7

8

9
class WikipediaReader(BasePydanticReader):
10
    """Wikipedia reader.
11

12
    Reads a page.
13

14
    """
15

16
    is_remote: bool = True
17

18
    def __init__(self) -> None:
19
        """Initialize with parameters."""
20
        try:
21
            import wikipedia  # noqa
22
        except ImportError:
23
            raise ImportError(
24
                "`wikipedia` package not found, please run `pip install wikipedia`"
25
            )
26

27
    @classmethod
28
    def class_name(cls) -> str:
29
        return "WikipediaReader"
30

31
    def load_data(self, pages: List[str], **load_kwargs: Any) -> List[Document]:
32
        """Load data from the input directory.
33

34
        Args:
35
            pages (List[str]): List of pages to read.
36

37
        """
38
        import wikipedia
39

40
        results = []
41
        for page in pages:
42
            wiki_page = wikipedia.page(page, **load_kwargs)
43
            page_content = wiki_page.content
44
            page_id = wiki_page.pageid
45
            results.append(Document(id_=page_id, text=page_content))
46
        return results
47