src/tools/wiki.py

from typing import Union

from langchain.docstore.base import Docstore
from langchain.docstore.document import Document



class Wiki(Docstore):
    """
    Wrapper around wikipedia API.
    """

    def __init__(self) -> None:
        """Check that wikipedia package is installed."""
        try:
            import wikipedia  # noqa: F401
        except ImportError:
            raise ValueError(
                "Could not import wikipedia python package. "
                "Please install it with `pip install wikipedia`."
            )

    @staticmethod
    def fetch(searched_page: str) -> Union[str, Document]:
        """
            Try to fetch for wiki page.

            If page exists, return the page summary, and a PageWithLookups object.
            If page does not exist, return similar entries.
        """
        import wikipedia

        try:
            # wikipedia.set_lang("fr")
            page_content = wikipedia.page(searched_page).content
            url = wikipedia.page(searched_page).url
            result: Union[str, Document] = Document(
                page_content=page_content, metadata={"page": url}
            )
        except wikipedia.PageError:
            result = f"Could not find [{searched_page}]. Similar: {wikipedia.search(searched_page)}"

        except wikipedia.DisambiguationError:
            result = f"Could not find [{searched_page}]. Similar: {wikipedia.search(searched_page)}"
        return result

    def search(searched_context: str) -> [str]:
        """
        Finds wiki page title in relation with the given context
        """
        import wikipedia

        try:
            # wikipedia.set_lang("fr")
            page_title_list = wikipedia.search(searched_context)
            result = page_title_list
        except wikipedia.PageError:
            result = f"Could not find [{searched_context}]."
        return result