diff --git a/langchain/utilities/wikipedia.py b/langchain/utilities/wikipedia.py index 6685ead8..259e8a44 100644 --- a/langchain/utilities/wikipedia.py +++ b/langchain/utilities/wikipedia.py @@ -39,7 +39,10 @@ class WikipediaAPIWrapper(BaseModel): """Run Wikipedia search and get page summaries.""" search_results = self.wiki_client.search(query) 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]) if summary is not None: summaries.append(summary) diff --git a/tests/integration_tests/test_wikipedia_api.py b/tests/integration_tests/test_wikipedia_api.py new file mode 100644 index 00000000..db82dde2 --- /dev/null +++ b/tests/integration_tests/test_wikipedia_api.py @@ -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