mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 15:48:29 +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 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
|
||||
|
||||
Reference in New Issue
Block a user