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:
akmhmgc 2023-04-03 06:00:27 +09:00 committed by GitHub
parent 715bd06f04
commit 94b2f536f3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 23 additions and 1 deletions

View File

@ -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)

View 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