diff --git a/agentstack/_tools/agentql/__init__.py b/agentstack/_tools/agentql/__init__.py index 469fcfcf..3b46d02e 100644 --- a/agentstack/_tools/agentql/__init__.py +++ b/agentstack/_tools/agentql/__init__.py @@ -8,7 +8,8 @@ API_KEY = os.getenv("AGENTQL_API_KEY") -def query_data(url: str, query: Optional[str], prompt: Optional[str]) -> dict: + +def extract_data(url: str, query: Optional[str], prompt: Optional[str]) -> dict: """ url: url of website to scrape query: described below @@ -44,37 +45,36 @@ def query_data(url: str, query: Optional[str], prompt: Optional[str]) -> dict: } ``` """ - payload = { - "url": url, - "query": query, - "prompt": prompt - } + payload = {"url": url, "query": query, "prompt": prompt} - headers = { - "X-API-Key": f"{API_KEY}", - "Content-Type": "application/json" - } + headers = {"X-API-Key": f"{API_KEY}", "Content-Type": "application/json"} try: response = httpx.post( QUERY_DATA_ENDPOINT, headers=headers, json=payload, - timeout=API_TIMEOUT_SECONDS + timeout=API_TIMEOUT_SECONDS, ) response.raise_for_status() except httpx.HTTPStatusError as e: response = e.response if response.status_code in [401, 403]: - raise ValueError("Please, provide a valid API Key. You can create one at https://dev.agentql.com.") from e + raise ValueError( + "Please, provide a valid API Key. You can create one at https://dev.agentql.com." + ) from e else: try: error_json = response.json() - msg = error_json["error_info"] if "error_info" in error_json else error_json["detail"] + msg = ( + error_json["error_info"] + if "error_info" in error_json + else error_json["detail"] + ) except (ValueError, TypeError): msg = f"HTTP {e}." raise ValueError(msg) from e else: json = response.json() - return json["data"] \ No newline at end of file + return json["data"] diff --git a/agentstack/_tools/agentql/config.json b/agentstack/_tools/agentql/config.json index 154be031..b16b5ed3 100644 --- a/agentstack/_tools/agentql/config.json +++ b/agentstack/_tools/agentql/config.json @@ -6,6 +6,6 @@ "env": { "AGENTQL_API_KEY": "..." }, - "tools": ["query_data"], + "tools": ["extract_data"], "cta": "Create your AgentQL API key at https://dev.agentql.com" }