Spaces:
Runtime error
Runtime error
File size: 1,531 Bytes
fb9c306 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
from dataclasses import dataclass
import requests
from fastapi import HTTPException
from graphgen.utils import logger
GOOGLE_SEARCH_ENDPOINT = "https://customsearch.googleapis.com/customsearch/v1"
@dataclass
class GoogleSearch:
def __init__(self, subscription_key: str, cx: str):
"""
Initialize the Google Search client with the subscription key and custom search engine ID.
:param subscription_key: Your Google API subscription key.
:param cx: Your custom search engine ID.
"""
self.subscription_key = subscription_key
self.cx = cx
def search(self, query: str, num_results: int = 1):
"""
Search with Google and return the contexts.
:param query: The search query.
:param num_results: The number of results to return.
:return: A list of search results.
"""
params = {
"key": self.subscription_key,
"cx": self.cx,
"q": query,
"num": num_results,
}
response = requests.get(GOOGLE_SEARCH_ENDPOINT, params=params, timeout=10)
if not response.ok:
logger.error("Search engine error: %s", response.text)
raise HTTPException(response.status_code, "Search engine error.")
json_content = response.json()
try:
contexts = json_content["items"][:num_results]
except KeyError:
logger.error("Error encountered: %s", json_content)
return []
return contexts
|