apply ruff formatting
This commit is contained in:
@@ -32,7 +32,7 @@ async def search_custom(
|
||||
date_restrict: Optional[str] = None,
|
||||
file_type: Optional[str] = None,
|
||||
language: Optional[str] = None,
|
||||
country: Optional[str] = None
|
||||
country: Optional[str] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Performs a search using Google Custom Search JSON API.
|
||||
@@ -55,54 +55,58 @@ async def search_custom(
|
||||
str: Formatted search results including title, link, and snippet for each result.
|
||||
"""
|
||||
# Get API key and search engine ID from environment
|
||||
api_key = os.environ.get('GOOGLE_PSE_API_KEY')
|
||||
api_key = os.environ.get("GOOGLE_PSE_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError("GOOGLE_PSE_API_KEY environment variable not set. Please set it to your Google Custom Search API key.")
|
||||
|
||||
cx = os.environ.get('GOOGLE_PSE_ENGINE_ID')
|
||||
if not cx:
|
||||
raise ValueError("GOOGLE_PSE_ENGINE_ID environment variable not set. Please set it to your Programmable Search Engine ID.")
|
||||
raise ValueError(
|
||||
"GOOGLE_PSE_API_KEY environment variable not set. Please set it to your Google Custom Search API key."
|
||||
)
|
||||
|
||||
logger.info(f"[search_custom] Invoked. Email: '{user_google_email}', Query: '{q}', CX: '{cx}'")
|
||||
cx = os.environ.get("GOOGLE_PSE_ENGINE_ID")
|
||||
if not cx:
|
||||
raise ValueError(
|
||||
"GOOGLE_PSE_ENGINE_ID environment variable not set. Please set it to your Programmable Search Engine ID."
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"[search_custom] Invoked. Email: '{user_google_email}', Query: '{q}', CX: '{cx}'"
|
||||
)
|
||||
|
||||
# Build the request parameters
|
||||
params = {
|
||||
'key': api_key,
|
||||
'cx': cx,
|
||||
'q': q,
|
||||
'num': num,
|
||||
'start': start,
|
||||
'safe': safe
|
||||
"key": api_key,
|
||||
"cx": cx,
|
||||
"q": q,
|
||||
"num": num,
|
||||
"start": start,
|
||||
"safe": safe,
|
||||
}
|
||||
|
||||
# Add optional parameters
|
||||
if search_type:
|
||||
params['searchType'] = search_type
|
||||
params["searchType"] = search_type
|
||||
if site_search:
|
||||
params['siteSearch'] = site_search
|
||||
params["siteSearch"] = site_search
|
||||
if site_search_filter:
|
||||
params['siteSearchFilter'] = site_search_filter
|
||||
params["siteSearchFilter"] = site_search_filter
|
||||
if date_restrict:
|
||||
params['dateRestrict'] = date_restrict
|
||||
params["dateRestrict"] = date_restrict
|
||||
if file_type:
|
||||
params['fileType'] = file_type
|
||||
params["fileType"] = file_type
|
||||
if language:
|
||||
params['lr'] = language
|
||||
params["lr"] = language
|
||||
if country:
|
||||
params['cr'] = country
|
||||
params["cr"] = country
|
||||
|
||||
# Execute the search request
|
||||
result = await asyncio.to_thread(
|
||||
service.cse().list(**params).execute
|
||||
)
|
||||
result = await asyncio.to_thread(service.cse().list(**params).execute)
|
||||
|
||||
# Extract search information
|
||||
search_info = result.get('searchInformation', {})
|
||||
total_results = search_info.get('totalResults', '0')
|
||||
search_time = search_info.get('searchTime', 0)
|
||||
search_info = result.get("searchInformation", {})
|
||||
total_results = search_info.get("totalResults", "0")
|
||||
search_time = search_info.get("searchTime", 0)
|
||||
|
||||
# Extract search results
|
||||
items = result.get('items', [])
|
||||
items = result.get("items", [])
|
||||
|
||||
# Format the response
|
||||
confirmation_message = f"""Search Results for {user_google_email}:
|
||||
@@ -117,43 +121,46 @@ async def search_custom(
|
||||
if items:
|
||||
confirmation_message += "Results:\n"
|
||||
for i, item in enumerate(items, start):
|
||||
title = item.get('title', 'No title')
|
||||
link = item.get('link', 'No link')
|
||||
snippet = item.get('snippet', 'No description available').replace('\n', ' ')
|
||||
title = item.get("title", "No title")
|
||||
link = item.get("link", "No link")
|
||||
snippet = item.get("snippet", "No description available").replace("\n", " ")
|
||||
|
||||
confirmation_message += f"\n{i}. {title}\n"
|
||||
confirmation_message += f" URL: {link}\n"
|
||||
confirmation_message += f" Snippet: {snippet}\n"
|
||||
|
||||
# Add additional metadata if available
|
||||
if 'pagemap' in item:
|
||||
pagemap = item['pagemap']
|
||||
if 'metatags' in pagemap and pagemap['metatags']:
|
||||
metatag = pagemap['metatags'][0]
|
||||
if 'og:type' in metatag:
|
||||
if "pagemap" in item:
|
||||
pagemap = item["pagemap"]
|
||||
if "metatags" in pagemap and pagemap["metatags"]:
|
||||
metatag = pagemap["metatags"][0]
|
||||
if "og:type" in metatag:
|
||||
confirmation_message += f" Type: {metatag['og:type']}\n"
|
||||
if 'article:published_time' in metatag:
|
||||
confirmation_message += f" Published: {metatag['article:published_time'][:10]}\n"
|
||||
if "article:published_time" in metatag:
|
||||
confirmation_message += (
|
||||
f" Published: {metatag['article:published_time'][:10]}\n"
|
||||
)
|
||||
else:
|
||||
confirmation_message += "\nNo results found."
|
||||
|
||||
# Add information about pagination
|
||||
queries = result.get('queries', {})
|
||||
if 'nextPage' in queries:
|
||||
next_start = queries['nextPage'][0].get('startIndex', 0)
|
||||
confirmation_message += f"\n\nTo see more results, search again with start={next_start}"
|
||||
queries = result.get("queries", {})
|
||||
if "nextPage" in queries:
|
||||
next_start = queries["nextPage"][0].get("startIndex", 0)
|
||||
confirmation_message += (
|
||||
f"\n\nTo see more results, search again with start={next_start}"
|
||||
)
|
||||
|
||||
logger.info(f"Search completed successfully for {user_google_email}")
|
||||
return confirmation_message
|
||||
|
||||
|
||||
@server.tool()
|
||||
@handle_http_errors("get_search_engine_info", is_read_only=True, service_type="customsearch")
|
||||
@handle_http_errors(
|
||||
"get_search_engine_info", is_read_only=True, service_type="customsearch"
|
||||
)
|
||||
@require_google_service("customsearch", "customsearch")
|
||||
async def get_search_engine_info(
|
||||
service,
|
||||
user_google_email: str
|
||||
) -> str:
|
||||
async def get_search_engine_info(service, user_google_email: str) -> str:
|
||||
"""
|
||||
Retrieves metadata about a Programmable Search Engine.
|
||||
|
||||
@@ -164,31 +171,35 @@ async def get_search_engine_info(
|
||||
str: Information about the search engine including its configuration and available refinements.
|
||||
"""
|
||||
# Get API key and search engine ID from environment
|
||||
api_key = os.environ.get('GOOGLE_PSE_API_KEY')
|
||||
api_key = os.environ.get("GOOGLE_PSE_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError("GOOGLE_PSE_API_KEY environment variable not set. Please set it to your Google Custom Search API key.")
|
||||
|
||||
cx = os.environ.get('GOOGLE_PSE_ENGINE_ID')
|
||||
if not cx:
|
||||
raise ValueError("GOOGLE_PSE_ENGINE_ID environment variable not set. Please set it to your Programmable Search Engine ID.")
|
||||
raise ValueError(
|
||||
"GOOGLE_PSE_API_KEY environment variable not set. Please set it to your Google Custom Search API key."
|
||||
)
|
||||
|
||||
logger.info(f"[get_search_engine_info] Invoked. Email: '{user_google_email}', CX: '{cx}'")
|
||||
cx = os.environ.get("GOOGLE_PSE_ENGINE_ID")
|
||||
if not cx:
|
||||
raise ValueError(
|
||||
"GOOGLE_PSE_ENGINE_ID environment variable not set. Please set it to your Programmable Search Engine ID."
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"[get_search_engine_info] Invoked. Email: '{user_google_email}', CX: '{cx}'"
|
||||
)
|
||||
|
||||
# Perform a minimal search to get the search engine context
|
||||
params = {
|
||||
'key': api_key,
|
||||
'cx': cx,
|
||||
'q': 'test', # Minimal query to get metadata
|
||||
'num': 1
|
||||
"key": api_key,
|
||||
"cx": cx,
|
||||
"q": "test", # Minimal query to get metadata
|
||||
"num": 1,
|
||||
}
|
||||
|
||||
result = await asyncio.to_thread(
|
||||
service.cse().list(**params).execute
|
||||
)
|
||||
result = await asyncio.to_thread(service.cse().list(**params).execute)
|
||||
|
||||
# Extract context information
|
||||
context = result.get('context', {})
|
||||
title = context.get('title', 'Unknown')
|
||||
context = result.get("context", {})
|
||||
title = context.get("title", "Unknown")
|
||||
|
||||
confirmation_message = f"""Search Engine Information for {user_google_email}:
|
||||
- Search Engine ID: {cx}
|
||||
@@ -196,18 +207,18 @@ async def get_search_engine_info(
|
||||
"""
|
||||
|
||||
# Add facet information if available
|
||||
if 'facets' in context:
|
||||
if "facets" in context:
|
||||
confirmation_message += "\nAvailable Refinements:\n"
|
||||
for facet in context['facets']:
|
||||
for facet in context["facets"]:
|
||||
for item in facet:
|
||||
label = item.get('label', 'Unknown')
|
||||
anchor = item.get('anchor', 'Unknown')
|
||||
label = item.get("label", "Unknown")
|
||||
anchor = item.get("anchor", "Unknown")
|
||||
confirmation_message += f" - {label} (anchor: {anchor})\n"
|
||||
|
||||
# Add search information
|
||||
search_info = result.get('searchInformation', {})
|
||||
search_info = result.get("searchInformation", {})
|
||||
if search_info:
|
||||
total_results = search_info.get('totalResults', 'Unknown')
|
||||
total_results = search_info.get("totalResults", "Unknown")
|
||||
confirmation_message += "\nSearch Statistics:\n"
|
||||
confirmation_message += f" - Total indexed results: {total_results}\n"
|
||||
|
||||
@@ -216,7 +227,9 @@ async def get_search_engine_info(
|
||||
|
||||
|
||||
@server.tool()
|
||||
@handle_http_errors("search_custom_siterestrict", is_read_only=True, service_type="customsearch")
|
||||
@handle_http_errors(
|
||||
"search_custom_siterestrict", is_read_only=True, service_type="customsearch"
|
||||
)
|
||||
@require_google_service("customsearch", "customsearch")
|
||||
async def search_custom_siterestrict(
|
||||
service,
|
||||
@@ -225,7 +238,7 @@ async def search_custom_siterestrict(
|
||||
sites: List[str],
|
||||
num: int = 10,
|
||||
start: int = 1,
|
||||
safe: Literal["active", "moderate", "off"] = "off"
|
||||
safe: Literal["active", "moderate", "off"] = "off",
|
||||
) -> str:
|
||||
"""
|
||||
Performs a search restricted to specific sites using Google Custom Search.
|
||||
@@ -241,7 +254,9 @@ async def search_custom_siterestrict(
|
||||
Returns:
|
||||
str: Formatted search results from the specified sites.
|
||||
"""
|
||||
logger.info(f"[search_custom_siterestrict] Invoked. Email: '{user_google_email}', Query: '{q}', Sites: {sites}")
|
||||
logger.info(
|
||||
f"[search_custom_siterestrict] Invoked. Email: '{user_google_email}', Query: '{q}', Sites: {sites}"
|
||||
)
|
||||
|
||||
# Build site restriction query
|
||||
site_query = " OR ".join([f"site:{site}" for site in sites])
|
||||
@@ -254,5 +269,5 @@ async def search_custom_siterestrict(
|
||||
q=full_query,
|
||||
num=num,
|
||||
start=start,
|
||||
safe=safe
|
||||
)
|
||||
safe=safe,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user