forked from Archives/langchain
Modify output for wikipedia api wrapper (#2287)
## Description Thanks for the quick maintenance for great repository!! I modified wikipedia api wrapper ## Details - Add output for missing search results - Add tests
This commit is contained in:
parent
715bd06f04
commit
94b2f536f3
@ -39,7 +39,10 @@ class WikipediaAPIWrapper(BaseModel):
|
|||||||
"""Run Wikipedia search and get page summaries."""
|
"""Run Wikipedia search and get page summaries."""
|
||||||
search_results = self.wiki_client.search(query)
|
search_results = self.wiki_client.search(query)
|
||||||
summaries = []
|
summaries = []
|
||||||
for i in range(min(self.top_k_results, len(search_results))):
|
len_search_results = len(search_results)
|
||||||
|
if len_search_results == 0:
|
||||||
|
return "No good Wikipedia Search Result was found"
|
||||||
|
for i in range(min(self.top_k_results, len_search_results)):
|
||||||
summary = self.fetch_formatted_page_summary(search_results[i])
|
summary = self.fetch_formatted_page_summary(search_results[i])
|
||||||
if summary is not None:
|
if summary is not None:
|
||||||
summaries.append(summary)
|
summaries.append(summary)
|
||||||
|
19
tests/integration_tests/test_wikipedia_api.py
Normal file
19
tests/integration_tests/test_wikipedia_api.py
Normal file
@ -0,0 +1,19 @@
|
|||||||
|
"""Integration test for Wikipedia API Wrapper."""
|
||||||
|
from langchain.utilities import WikipediaAPIWrapper
|
||||||
|
|
||||||
|
|
||||||
|
def test_call() -> None:
|
||||||
|
"""Test that WikipediaAPIWrapper returns correct answer"""
|
||||||
|
|
||||||
|
wikipedia = WikipediaAPIWrapper()
|
||||||
|
output = wikipedia.run("HUNTER X HUNTER")
|
||||||
|
assert "Yoshihiro Togashi" in output
|
||||||
|
|
||||||
|
|
||||||
|
def test_no_result_call() -> None:
|
||||||
|
"""Test that call gives no result."""
|
||||||
|
wikipedia = WikipediaAPIWrapper()
|
||||||
|
output = wikipedia.run(
|
||||||
|
"NORESULTCALL_NORESULTCALL_NORESULTCALL_NORESULTCALL_NORESULTCALL_NORESULTCALL"
|
||||||
|
)
|
||||||
|
assert "No good Wikipedia Search Result was found" == output
|
Loading…
Reference in New Issue
Block a user