mirror of
https://gitee.com/dify_ai/dify.git
synced 2024-12-02 11:18:19 +08:00
52 lines
2.1 KiB
Python
52 lines
2.1 KiB
Python
from langchain import SerpAPIWrapper
|
|
from pydantic import Field, BaseModel
|
|
|
|
|
|
class OptimizedSerpAPIInput(BaseModel):
|
|
query: str = Field(..., description="search query.")
|
|
|
|
|
|
class OptimizedSerpAPIWrapper(SerpAPIWrapper):
|
|
|
|
@staticmethod
|
|
def _process_response(res: dict, num_results: int = 5) -> str:
|
|
"""Process response from SerpAPI."""
|
|
if "error" in res.keys():
|
|
raise ValueError(f"Got error from SerpAPI: {res['error']}")
|
|
if "answer_box" in res.keys() and type(res["answer_box"]) == list:
|
|
res["answer_box"] = res["answer_box"][0]
|
|
if "answer_box" in res.keys() and "answer" in res["answer_box"].keys():
|
|
toret = res["answer_box"]["answer"]
|
|
elif "answer_box" in res.keys() and "snippet" in res["answer_box"].keys():
|
|
toret = res["answer_box"]["snippet"]
|
|
elif (
|
|
"answer_box" in res.keys()
|
|
and "snippet_highlighted_words" in res["answer_box"].keys()
|
|
):
|
|
toret = res["answer_box"]["snippet_highlighted_words"][0]
|
|
elif (
|
|
"sports_results" in res.keys()
|
|
and "game_spotlight" in res["sports_results"].keys()
|
|
):
|
|
toret = res["sports_results"]["game_spotlight"]
|
|
elif (
|
|
"shopping_results" in res.keys()
|
|
and "title" in res["shopping_results"][0].keys()
|
|
):
|
|
toret = res["shopping_results"][:3]
|
|
elif (
|
|
"knowledge_graph" in res.keys()
|
|
and "description" in res["knowledge_graph"].keys()
|
|
):
|
|
toret = res["knowledge_graph"]["description"]
|
|
elif 'organic_results' in res.keys() and len(res['organic_results']) > 0:
|
|
toret = ""
|
|
for result in res["organic_results"][:num_results]:
|
|
if "link" in result:
|
|
toret += "----------------\nlink: " + result["link"] + "\n"
|
|
if "snippet" in result:
|
|
toret += "snippet: " + result["snippet"] + "\n"
|
|
else:
|
|
toret = "No good search result found"
|
|
return "search result:\n" + toret
|