| 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 |
| 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: |
| |
| 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: |
| |
| page_title_list = wikipedia.search(searched_context) |
| result = page_title_list |
| except wikipedia.PageError: |
| result = f"Could not find [{searched_context}]." |
| return result |
|
|
|
|
|
|
