|
|
from typing import Optional, Union |
|
|
from langchain.docstore.document import Document |
|
|
import wikipedia |
|
|
|
|
|
class LangChainWiki: |
|
|
|
|
|
def __init__(self) -> None: |
|
|
self.document: Optional[Document] = None |
|
|
self.lookup_str = "" |
|
|
self.lookup_index = 0 |
|
|
|
|
|
def search(self, search: str) -> Union[str, Document]: |
|
|
def _try_search(term: str) -> Union[str, Document]: |
|
|
try: |
|
|
page_content = wikipedia.page(search).content |
|
|
url = wikipedia.page(search).url |
|
|
result: Union[str, Document] = Document( page_content=page_content, metadata={"page": url} ) |
|
|
return result |
|
|
except wikipedia.PageError: |
|
|
return f"Could not find [{term}]. Similar: {wikipedia.search(term)}" |
|
|
except wikipedia.DisambiguationError: |
|
|
return f"Could not find [{term}]. Similar: {wikipedia.search(term)}" |
|
|
except Exception: |
|
|
return f"Could not find [{term}]. Similar: {wikipedia.search(term)}" |
|
|
|
|
|
result = _try_search(search) |
|
|
|
|
|
if isinstance(result, str) and "Similar:" in result: |
|
|
try: |
|
|
similar = wikipedia.search(search) |
|
|
if similar: |
|
|
fallback = similar[0] |
|
|
print(f"[INFO] Falling back to similar term: {fallback}") |
|
|
result = _try_search(fallback) |
|
|
except Exception as e: |
|
|
print(f"[ERROR] Could not fetch similar terms: {e}") |
|
|
|
|
|
if isinstance(result, Document): |
|
|
self.document = result |
|
|
return self._sumary |
|
|
else: |
|
|
self.document = None |
|
|
return result |
|
|
|
|
|
def lookup(self, term: str): |
|
|
if self.document is None: |
|
|
raise ValueError("Cannot lookup without a successful search first") |
|
|
if term.lower() != self.lookup_str: |
|
|
self.lookup_str = term.lower() |
|
|
self.lookup_index = 0 |
|
|
else: |
|
|
self.lookup_index += 1 |
|
|
lookups = [p for p in self._paragraphs if self.lookup_str in p.lower()] |
|
|
if len(lookups) == 0: |
|
|
return "No Results" |
|
|
elif self.lookup_index >= len(lookups): |
|
|
return "No More Results" |
|
|
else: |
|
|
result_prefix = f"(Result {self.lookup_index + 1}/{len(lookups)})" |
|
|
return f"{result_prefix} {lookups[self.lookup_index]}" |
|
|
|
|
|
@property |
|
|
def _sumary(self) -> str: |
|
|
return self._paragraphs[0] |
|
|
|
|
|
@property |
|
|
def _paragraphs(self) -> list[str]: |
|
|
if self.document is None: |
|
|
raise ValueError("Cannot get paragraphs without a document") |
|
|
return self.document.page_content.split("\n\n") |
|
|
|
