Update serper_dev_tool.py

Added two additional functionalities:

1) added the ability to save the server results to a file
2) added the ability to  set the number of results returned

Can be used as follows:

serper_tool = SerperDevTool(file_save=True, n_results=20)
This commit is contained in:
Rip&Tear
2024-06-13 12:53:35 +08:00
committed by GitHub
parent 53c7d815ae
commit 5e8e711170

View File

@@ -6,6 +6,14 @@ from typing import Type, Any
from pydantic.v1 import BaseModel, Field
from crewai_tools.tools.base_tool import BaseTool
def _save_results_to_file(content: str) -> None:
"""Saves the search results to a file."""
filename = f"search_results_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.txt"
with open(filename, 'w') as file:
file.write(content)
print(f"Results saved to {filename}")
class SerperDevToolSchema(BaseModel):
"""Input for SerperDevTool."""
search_query: str = Field(..., description="Mandatory search query you want to use to search the internet")
@@ -15,17 +23,22 @@ class SerperDevTool(BaseTool):
description: str = "A tool that can be used to search the internet with a search_query."
args_schema: Type[BaseModel] = SerperDevToolSchema
search_url: str = "https://google.serper.dev/search"
n_results: int = 10
n_results: int = Field(default=10, description="Number of search results to return")
save_file: bool = Field(default=False, description="Flag to determine whether to save the results to a file")
def _run(
self,
**kwargs: Any,
) -> Any:
save_file = kwargs.get('save_file', self.save_file)
n_results = kwargs.get('n_results', self.n_results)
search_query = kwargs.get('search_query')
if search_query is None:
search_query = kwargs.get('query')
payload = json.dumps({"q": search_query})
payload = json.dumps({"q": search_query, "num": n_results})
headers = {
'X-API-KEY': os.environ['SERPER_API_KEY'],
'content-type': 'application/json'
@@ -47,6 +60,8 @@ class SerperDevTool(BaseTool):
next
content = '\n'.join(string)
if save_file:
_save_results_to_file(content)
return f"\nSearch results: {content}\n"
else:
return results