prevent DuckDuckGoSearchAPIWrapper from consuming top result (#6727)

remove the `next` call that checks for None on the results generator
This commit is contained in:
Gabriel Altay 2023-06-25 22:54:15 -04:00 committed by GitHub
parent 87802c86d9
commit f9771700e4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -49,14 +49,15 @@ class DuckDuckGoSearchAPIWrapper(BaseModel):
safesearch=self.safesearch, safesearch=self.safesearch,
timelimit=self.time, timelimit=self.time,
) )
if results is None or next(results, None) is None: if results is None:
return ["No good DuckDuckGo Search Result was found"] return ["No good DuckDuckGo Search Result was found"]
snippets = [] snippets = []
for i, res in enumerate(results, 1): for i, res in enumerate(results, 1):
snippets.append(res["body"]) if res is not None:
if i == self.max_results: snippets.append(res["body"])
if len(snippets) == self.max_results:
break break
return snippets return snippets
def run(self, query: str) -> str: def run(self, query: str) -> str:
snippets = self.get_snippets(query) snippets = self.get_snippets(query)
@ -84,7 +85,7 @@ class DuckDuckGoSearchAPIWrapper(BaseModel):
safesearch=self.safesearch, safesearch=self.safesearch,
timelimit=self.time, timelimit=self.time,
) )
if results is None or next(results, None) is None: if results is None:
return [{"Result": "No good DuckDuckGo Search Result was found"}] return [{"Result": "No good DuckDuckGo Search Result was found"}]
def to_metadata(result: Dict) -> Dict[str, str]: def to_metadata(result: Dict) -> Dict[str, str]:
@ -96,7 +97,8 @@ class DuckDuckGoSearchAPIWrapper(BaseModel):
formatted_results = [] formatted_results = []
for i, res in enumerate(results, 1): for i, res in enumerate(results, 1):
formatted_results.append(to_metadata(res)) if res is not None:
if i == num_results: formatted_results.append(to_metadata(res))
if len(formatted_results) == num_results:
break break
return formatted_results return formatted_results