""" title: Bigquery toolkit author: Wendy description: A tool for interacting with Bigquery - query database, describe table and list tables. required_open_webui_version: 0.4.0 requirements: google-cloud-bigquery version: 0.4.0 licence: MIT """ import json import logging import os from google.cloud import bigquery import google.auth from typing import Awaitable, Callable, Dict, Any from pydantic import BaseModel, Field def setup_logger(): name = "GoogleTools" logger = logging.getLogger(name) if not logger.handlers: logger.setLevel(logging.DEBUG) handler = logging.StreamHandler() handler.set_name(name) formatter = logging.Formatter( "%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) handler.setFormatter(formatter) logger.addHandler(handler) logger.propagate = False return logger logger = setup_logger() def beautify(rows, column_names): """ Create a beautiful ASCII table from SQL query results. Args: rows: List of dictionaries, each containing the values for a row. column_names: List of column names. Returns: A formatted string representing the table. """ if not rows or not column_names: return "No data to display." # Convert all values to strings and find the maximum width for each column str_rows = [] col_widths = {col: len(col) for col in column_names} for row in rows: str_row = {} for col in column_names: value = str(row.get(col, "")) str_row[col] = value col_widths[col] = max(col_widths[col], len(value)) str_rows.append(str_row) # Calculate the total width of the table total_width = sum(col_widths.values()) + (3 * len(column_names)) + 1 # Create the table result = [] # Top border result.append("+" + "-" * (total_width - 2) + "+") # Header row header = "|" for col in column_names: header += f" {col.ljust(col_widths[col])} |" result.append(header) # Header separator separator = "+" for col in column_names: separator += "-" * (col_widths[col] + 2) + "+" result.append(separator) # Data rows for row in str_rows: data_row = "|" for col in column_names: data_row += f" {row[col].ljust(col_widths[col])} |" result.append(data_row) # Bottom border result.append("+" + "-" * (total_width - 2) + "+") return "\n".join(result) class EventEmitter: def __init__(self, event_emitter: Callable[[dict], Awaitable[None]]): self.event_emitter = event_emitter pass async def emit_status(self, description: str, done: bool, error: bool = False): await self.event_emitter( { "data": { "description": f"{done and (error and '❌' or '✅') or '🔎'} {description}", "status": done and "complete" or "in_progress", "done": done, }, "type": "status", } ) async def emit_message(self, content: str): await self.event_emitter({"data": {"content": content}, "type": "message"}) async def emit_query_result(self, query: str, query_result: list[dict] | list[str]): await self.event_emitter( { "type": "citation", "data": {"query": query, "result": query_result}, } ) class BigQueryDatabase: def __init__( self, project: str, datasets_filter: list[str], credentials=None, location: str = None, ): """Initialize a BigQuery database client""" if not project: raise ValueError("Project is required") # if not location: # raise ValueError("Location is required") if credentials is None: self.client = bigquery.Client(project=project, location=location) else: self.client = bigquery.Client(project=project, credentials=credentials) self.datasets_filter = datasets_filter def execute_query(self, query: str, params: dict[str, Any] | None = None) -> str: """Execute a read query and return the results in CSV format""" if params: job = self.client.query( query, job_config=bigquery.QueryJobConfig(query_parameters=params) ) else: job = self.client.query(query) results = job.result() rows = [dict(row.items()) for row in results] column_names = [s.name for s in results.schema] return rows, column_names def list_tables(self) -> list[str]: # bug: too much table!!! """List all tables in the BigQuery database""" if self.datasets_filter: datasets = [ self.client.dataset(dataset) for dataset in self.datasets_filter ] else: datasets = list(self.client.list_datasets()) tables = [] for dataset in datasets: dataset_tables = self.client.list_tables(dataset.dataset_id) tables.extend( [f"{dataset.dataset_id}.{table.table_id}" for table in dataset_tables] ) return tables def describe_table(self, table_name: str) -> list[dict[str, Any]]: """Describe a table in the BigQuery database""" parts = table_name.split(".") if len(parts) != 2: raise ValueError(f"Invalid table name: {table_name}") dataset_id = parts[0] table_id = parts[1] query = f""" SELECT ddl FROM {dataset_id}.INFORMATION_SCHEMA.TABLES WHERE table_name = @table_name; """ result, column_names = self.execute_query( query, params=[ bigquery.ScalarQueryParameter("table_name", "STRING", table_id), ], ) return result, column_names class Tools: def __init__(self): self.valves = self.Valves() # self.citation = False class Valves(BaseModel): project: str = Field( "", description="Your Google Cloud Project name where the BigQuery resources are located.", ) location: str = Field( "US", description="The geographic location of your BigQuery datasets. Default is 'US'.", ) bigquery_dataset: str = Field( "", description="The BigQuery datasets visible to the LLM. Multiple datasets should be separated by commas. An empty string ('') means all datasets are visible.", ) def get_google_creds(self): credentials, project = google.auth.default( scopes=[ "https://www.googleapis.com/auth/drive", "https://www.googleapis.com/auth/cloud-platform", ] ) return credentials async def execute_query( self, query: str, __event_emitter__: Callable[[dict], Awaitable[None]], __user__: dict = {}, ) -> str: """ Execute a SELECT query on the Bigquery database. :param query: SELECT SQL query to execute using BigQuery dialect. :return: A table format query result """ event_emitter = EventEmitter(__event_emitter__) await event_emitter.emit_status(f"execute query: {query}", False) try: datasets_filter = self.valves.bigquery_dataset.split(",") cred = self.get_google_creds() # db = BigQueryDatabase( # self.valves.project, self.valves.location, datasets_filter # ) db = BigQueryDatabase( credentials=cred, project=self.valves.project, datasets_filter=datasets_filter, ) logger.info(query) rows, column_names = db.execute_query(query) table = beautify(rows=rows, column_names=column_names) return_str = f"Query executed successfully. Below is the actual result of the query `{query}` running against the Bigquery database:\n\n" return_str += table await event_emitter.emit_status(f"Done execute query: {query}", True) # await event_emitter.emit_query_result( # query=query, query_result=json.dumps(results) # ) logger.info(return_str) return return_str except Exception as e: logger.error(e) await event_emitter.emit_status( f"Failed exce query {query}: {e}.", True, True ) return str(e) async def list_tables( self, __event_emitter__: Callable[[dict], Awaitable[None]], __user__: dict = {}, ) -> str: """ List all tables in the Bigquery database Note: This returns a list of pages that match the search query. :return: A list of table name. If no table are found, an empty list is returned. """ event_emitter = EventEmitter(__event_emitter__) await event_emitter.emit_status(f"Listing tables...", False) try: datasets_filter = self.valves.bigquery_dataset.split(",") cred = self.get_google_creds() # db = BigQueryDatabase( # self.valves.project, self.valves.location, datasets_filter # ) db = BigQueryDatabase( credentials=cred, project=self.valves.project, datasets_filter=datasets_filter, ) tables = db.list_tables() await event_emitter.emit_status(f"Done list tables.", True) # await event_emitter.emit_query_result( # query="list table", query_result=json.dumps(tables) # ) logger.info(tables) return json.dumps(tables) except Exception as e: await event_emitter.emit_status(f"Failed list tables: {e}.", True, True) logger.error(e) return str(e) async def describe_table( self, table_name: str, __event_emitter__: Callable[[dict], Awaitable[None]], __user__: dict = {}, ) -> str: """ Get the schema information for a specific table at bigquery database. Note: This returns a list of pages that match the search query. :param table_name: Name of the table to describe (e.g. my_dataset.my_table). :return: A list of search results from bigquery in JSON format (id, title, body, link). If no results are found, an empty list is returned. """ event_emitter = EventEmitter(__event_emitter__) await event_emitter.emit_status(f"describe table: {table_name}", False) try: datasets_filter = self.valves.bigquery_dataset.split(",") cred = self.get_google_creds() # db = BigQueryDatabase( # self.valves.project, self.valves.location, datasets_filter # ) db = BigQueryDatabase( credentials=cred, project=self.valves.project, datasets_filter=datasets_filter, ) results, col_names = db.describe_table(table_name) await event_emitter.emit_status(f"Done describe table.", True) # await event_emitter.emit_query_result( # query=table_name, query_result=json.dumps(results) # ) logger.info(results[0][col_names]) logger.info(col_names) return json.dumps(results) except Exception as e: logger.error(e) await event_emitter.emit_status(f"Failed describe table: {e}.", True, True) return str(e)