mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 00:28:31 +00:00
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:
@@ -6,6 +6,14 @@ from typing import Type, Any
|
|||||||
from pydantic.v1 import BaseModel, Field
|
from pydantic.v1 import BaseModel, Field
|
||||||
from crewai_tools.tools.base_tool import BaseTool
|
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):
|
class SerperDevToolSchema(BaseModel):
|
||||||
"""Input for SerperDevTool."""
|
"""Input for SerperDevTool."""
|
||||||
search_query: str = Field(..., description="Mandatory search query you want to use to search the internet")
|
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."
|
description: str = "A tool that can be used to search the internet with a search_query."
|
||||||
args_schema: Type[BaseModel] = SerperDevToolSchema
|
args_schema: Type[BaseModel] = SerperDevToolSchema
|
||||||
search_url: str = "https://google.serper.dev/search"
|
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(
|
def _run(
|
||||||
self,
|
self,
|
||||||
**kwargs: Any,
|
**kwargs: Any,
|
||||||
) -> 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')
|
search_query = kwargs.get('search_query')
|
||||||
if search_query is None:
|
if search_query is None:
|
||||||
search_query = kwargs.get('query')
|
search_query = kwargs.get('query')
|
||||||
|
|
||||||
payload = json.dumps({"q": search_query})
|
payload = json.dumps({"q": search_query, "num": n_results})
|
||||||
headers = {
|
headers = {
|
||||||
'X-API-KEY': os.environ['SERPER_API_KEY'],
|
'X-API-KEY': os.environ['SERPER_API_KEY'],
|
||||||
'content-type': 'application/json'
|
'content-type': 'application/json'
|
||||||
@@ -47,6 +60,8 @@ class SerperDevTool(BaseTool):
|
|||||||
next
|
next
|
||||||
|
|
||||||
content = '\n'.join(string)
|
content = '\n'.join(string)
|
||||||
|
if save_file:
|
||||||
|
_save_results_to_file(content)
|
||||||
return f"\nSearch results: {content}\n"
|
return f"\nSearch results: {content}\n"
|
||||||
else:
|
else:
|
||||||
return results
|
return results
|
||||||
|
|||||||
Reference in New Issue
Block a user