"""
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()
