""" title: SQLite Access Agent authors: yage description: This tool provides access to a given SQLite database, whose path is supplied in the Valve. author_url: https://github.com/grapeot/ funding_url: https://github.com/open-webui version: 0.0.1 required_open_webui_version: 0.3.17 license: MIT """ import os import sqlite3 import csv from datetime import datetime from pydantic import BaseModel, Field """ This tool provides access to a given SQLite database. For the ease of file exchange, we choose to supply a Web URL as the database file path. Note this URL could be a local/intranet URL. It supports listing all tables, describing a table, and execute a read query against the sqlite database. """ class Tools: class Valves(BaseModel): db_url: str = Field( default="https://yage.ai/genai/example1/events.db", description="The URL of the SQLite database. We will try to download it to local.", ) db_local_path: str = Field( default="/tmp/sqlite.db", description="The local path of the SQLite database.", ) def __init__(self): """ Initialize the Tools class with the path to the SQLite database. If the database file does not exist at db_path, it will be downloaded from db_url. :param db_path: The path to the SQLite database file. """ print("Initializing SQLite tool class") self.citation = True self.valves = Tools.Valves() def download_db(self): if not os.path.exists(self.valves.db_local_path): import requests print( f"Downloading database from {self.valves.db_url} to {self.valves.db_local_path}" ) response = requests.get(self.valves.db_url) with open(self.valves.db_local_path, "wb") as f: f.write(response.content) def list_all_tables(self) -> str: """ You have access to a few SQLite tables in a SQLite database. The user might ask you about what tables you have access to. Invoke this function to list all tables in the SQLite database. :return: A string containing the names of all tables. """ print(f"Listing all tables in {self.valves.db_local_path}") self.download_db() conn = sqlite3.connect(self.valves.db_local_path) cursor = conn.cursor() cursor.execute("SELECT name FROM sqlite_master WHERE type='table';") tables = [row[0] for row in cursor.fetchall()] conn.close() print(tables) if tables: return ( "Here is a list of all the tables in the SQLite database:\n\n" + "\n".join(tables) ) else: return "No tables found." def table_data_schema(self, table_name: str) -> str: """ You have access to a few SQLite tables. The user might ask you more details a given table. Invoke this function to get the data schema of a specific table in the SQLite database, e.g. for data schemas. :param table_name: The name of the table to describe. :return: A string describing the data schema of the table. """ print(f"Describing table: {table_name}") self.download_db() conn = sqlite3.connect(self.valves.db_local_path) cursor = conn.cursor() cursor.execute(f"PRAGMA table_info({table_name});") columns = cursor.fetchall() conn.close() if not columns: print(f"No such table: {table_name}") return f"No such table: {table_name}" description = ( f"Table '{table_name}' in the SQLite database has the following columns:\n" ) for column in columns: description += f"- {column[1]} ({column[2]})\n" print(description) return description def execute_read_query(self, query: str) -> str: """ You have access to a few SQLite tables. The user might ask you to execute some SQL query. Invoke this function to execute a read query and get the result in csv format. :param query: The SQL query to execute. :return: A string containing the result of the query in csv format. """ print(f"Executing query: {query}") self.download_db() conn = sqlite3.connect(self.valves.db_local_path) cursor = conn.cursor() try: cursor.execute(query) rows = cursor.fetchall() if not rows: return "No data returned from query." # Fetch column names column_names = [description[0] for description in cursor.description] csv_data = f"Query executed successfully. Below is the actual result of the query {query} running against the database in csv format:\n\n" # Add column names to CSV data csv_data += ",".join(column_names) + "\n" for row in rows: csv_data += ",".join(map(str, row)) + "\n" print(csv_data) return csv_data except sqlite3.Error as e: return f"Error executing query: {str(e)}" finally: conn.close()