mirror of
https://github.com/aljazceru/Auto-GPT.git
synced 2026-01-04 06:44:31 +01:00
142 lines
3.9 KiB
Python
142 lines
3.9 KiB
Python
"""Google search command for Autogpt."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import time
|
|
from itertools import islice
|
|
|
|
from duckduckgo_search import DDGS
|
|
|
|
from autogpt.agent.agent import Agent
|
|
from autogpt.command_decorator import command
|
|
|
|
DUCKDUCKGO_MAX_ATTEMPTS = 3
|
|
|
|
|
|
@command(
|
|
"web_search",
|
|
"Searches the web",
|
|
{
|
|
"query": {
|
|
"type": "string",
|
|
"description": "The search query",
|
|
"required": True,
|
|
}
|
|
},
|
|
)
|
|
def web_search(query: str, agent: Agent, num_results: int = 8) -> str:
|
|
"""Return the results of a Google search
|
|
|
|
Args:
|
|
query (str): The search query.
|
|
num_results (int): The number of results to return.
|
|
|
|
Returns:
|
|
str: The results of the search.
|
|
"""
|
|
search_results = []
|
|
attempts = 0
|
|
|
|
while attempts < DUCKDUCKGO_MAX_ATTEMPTS:
|
|
if not query:
|
|
return json.dumps(search_results)
|
|
|
|
results = DDGS().text(query)
|
|
search_results = list(islice(results, num_results))
|
|
|
|
if search_results:
|
|
break
|
|
|
|
time.sleep(1)
|
|
attempts += 1
|
|
|
|
results = json.dumps(search_results, ensure_ascii=False, indent=4)
|
|
return safe_google_results(results)
|
|
|
|
|
|
@command(
|
|
"google",
|
|
"Google Search",
|
|
{
|
|
"query": {
|
|
"type": "string",
|
|
"description": "The search query",
|
|
"required": True,
|
|
}
|
|
},
|
|
lambda config: bool(config.google_api_key)
|
|
and bool(config.google_custom_search_engine_id),
|
|
"Configure google_api_key and custom_search_engine_id.",
|
|
)
|
|
def google(query: str, agent: Agent, num_results: int = 8) -> str | list[str]:
|
|
"""Return the results of a Google search using the official Google API
|
|
|
|
Args:
|
|
query (str): The search query.
|
|
num_results (int): The number of results to return.
|
|
|
|
Returns:
|
|
str: The results of the search.
|
|
"""
|
|
|
|
from googleapiclient.discovery import build
|
|
from googleapiclient.errors import HttpError
|
|
|
|
try:
|
|
# Get the Google API key and Custom Search Engine ID from the config file
|
|
api_key = agent.config.google_api_key
|
|
custom_search_engine_id = agent.config.google_custom_search_engine_id
|
|
|
|
# Initialize the Custom Search API service
|
|
service = build("customsearch", "v1", developerKey=api_key)
|
|
|
|
# Send the search query and retrieve the results
|
|
result = (
|
|
service.cse()
|
|
.list(q=query, cx=custom_search_engine_id, num=num_results)
|
|
.execute()
|
|
)
|
|
|
|
# Extract the search result items from the response
|
|
search_results = result.get("items", [])
|
|
|
|
# Create a list of only the URLs from the search results
|
|
search_results_links = [item["link"] for item in search_results]
|
|
|
|
except HttpError as e:
|
|
# Handle errors in the API call
|
|
error_details = json.loads(e.content.decode())
|
|
|
|
# Check if the error is related to an invalid or missing API key
|
|
if error_details.get("error", {}).get(
|
|
"code"
|
|
) == 403 and "invalid API key" in error_details.get("error", {}).get(
|
|
"message", ""
|
|
):
|
|
return "Error: The provided Google API key is invalid or missing."
|
|
else:
|
|
return f"Error: {e}"
|
|
# google_result can be a list or a string depending on the search results
|
|
|
|
# Return the list of search result URLs
|
|
return safe_google_results(search_results_links)
|
|
|
|
|
|
def safe_google_results(results: str | list) -> str:
|
|
"""
|
|
Return the results of a google search in a safe format.
|
|
|
|
Args:
|
|
results (str | list): The search results.
|
|
|
|
Returns:
|
|
str: The results of the search.
|
|
"""
|
|
if isinstance(results, list):
|
|
safe_message = json.dumps(
|
|
[result.encode("utf-8", "ignore").decode("utf-8") for result in results]
|
|
)
|
|
else:
|
|
safe_message = results.encode("utf-8", "ignore").decode("utf-8")
|
|
return safe_message
|