apply ruff formatting

This commit is contained in:
Taylor Wilsdon
2025-12-13 13:49:28 -08:00
parent 1d80a24ca4
commit 6b8352a354
50 changed files with 4010 additions and 2842 deletions

View File

@@ -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,
)