app.py

import gradio as gr
import pandas as pd
import sys
import os

from utils.tools import get_kb, get_schema, get_tables, get_meaning, execute_sqlite_query

@gr.mcp.tool()
def get_all_databases() -> list:
    """
    Get all available database names from the schema file.
    
    This function reads the database schema CSV file and extracts unique database names.
    
    Returns:
        list: A sorted list of unique database names available in the system.
    
    Example:
        >>> databases = get_all_databases()
        >>> print(databases)
        ['db1', 'db2', 'db3']
    """
    # 從 schema_df 中獲取所有唯一的 db_name
    schema_df = pd.read_csv("./data/db_schema.csv")
    return sorted(schema_df['db_name'].unique().tolist())

def kb_query(db_name, knowledge_keyword):
    """
    Query the knowledge base for a specific database with optional keyword filtering.
    
    This function retrieves knowledge base information for a specified database.
    If a keyword is provided, it filters the results based on that keyword.
    
    Args:
        db_name (str): The name of the database to query. Must not be empty.
        knowledge_keyword (str): Optional keyword to filter knowledge base results.
                               If empty or None, returns all knowledge for the database.
    
    Returns:
        pandas.DataFrame: Query results containing knowledge base data, or error message
                         if no database is selected or no results found.
    
    Example:
        >>> result = kb_query("sales_db", "customer")
        >>> print(result)
        # Returns DataFrame with customer-related knowledge from sales_db
    """
    if not db_name:
        return pd.DataFrame({"message": ["請先選擇資料庫"]})
    
    if not knowledge_keyword:
        result = get_kb(db_name)
    else:
        result = get_kb(db_name, knowledge_keyword)
    
    if len(result) == 0:
        return pd.DataFrame({"message": ["沒有找到相關知識"]})
    
    return result

def schema_query(db_name, table_name):
    """
    Query the schema structure for a specific table in a database.
    
    This function retrieves detailed schema information for a specified table
    within a given database, including column definitions, data types, and constraints.
    
    Args:
        db_name (str): The name of the database containing the table. Must not be empty.
        table_name (str): The name of the table to query schema for. Must not be empty.
    
    Returns:
        pandas.DataFrame: Query results containing table schema information, or error message
                         if parameters are missing or no schema found.
    
    Example:
        >>> result = schema_query("sales_db", "customers")
        >>> print(result)
        # Returns DataFrame with column definitions for customers table
    """
    if not db_name or not table_name:
        return pd.DataFrame({"message": ["請選擇資料庫和資料表"]})
    
    result = get_schema(db_name, table_name)
    
    if len(result) == 0:
        return pd.DataFrame({"message": ["沒有找到相關資料表結構"]})
    
    return result

def tables_query(db_name):
    """
    Get list of all tables available in a specific database.
    
    This function retrieves all table names that exist within the specified database.
    
    Args:
        db_name (str): The name of the database to query tables from. 
                      If empty or None, returns empty list.
    
    Returns:
        list: List of table names in the specified database. Returns empty list
              if database name is not provided or no tables found.
    
    Example:
        >>> tables = tables_query("sales_db")
        >>> print(tables)
        ['customers', 'orders', 'products', 'inventory']
    """
    if not db_name:
        return []
    
    return get_tables(db_name)

def meaning_query(db_name, table_name):
    """
    Query the meaning and description of columns in a specific table.
    
    This function retrieves detailed explanations and meanings for each column
    in the specified table, helping users understand the purpose and content
    of each field.
    
    Args:
        db_name (str): The name of the database containing the table. Must not be empty.
        table_name (str): The name of the table to query column meanings for. Must not be empty.
    
    Returns:
        pandas.DataFrame: Query results containing column meanings and descriptions,
                         or error message if parameters are missing or no meanings found.
    
    Example:
        >>> result = meaning_query("sales_db", "customers")
        >>> print(result)
        # Returns DataFrame with explanations for each column in customers table
    """
    if not db_name or not table_name:
        return pd.DataFrame({"message": ["請選擇資料庫和資料表"]})
    
    result = get_meaning(db_name, table_name)
    
    if len(result) == 0:
        return pd.DataFrame({"message": ["沒有找到相關欄位意義"]})
    
    return pd.DataFrame(result)

def execute_sqlite(db_name: str, query: str):
    """
    Execute an SQL query on a specified SQLite database and return results.

    This function allows direct SQL query execution on the selected database,
    returning either the query results or error messages in a standardized format.

    Args:
        db_name (str): The name of the database to execute the query against. Must not be empty.
        query (str): The SQL query to execute. Must not be empty.

    Returns:
        pandas.DataFrame: If query executes successfully, returns DataFrame with query results.
                         If query fails or parameters are missing, returns error message in DataFrame.

    Example:
        >>> result = execute_sqlite("sales_db", "SELECT * FROM customers LIMIT 10")
        >>> print(result)
        # Returns DataFrame with first 10 rows from customers table
    """
    if not db_name or not query:
        return pd.DataFrame({"message": ["請選擇資料庫和SQL"]})
    
    result = execute_sqlite_query(db_name, query)
    
    if len(result) == 0:
        return pd.DataFrame({"message": ["沒有找到相關資料"]})
    
    if result['success']:
        return pd.DataFrame(result['data'], columns=result['columns'])
    else:
        return pd.DataFrame({"message": [result['error']]})


# 建立 Gradio 界面
with gr.Blocks(title="資料庫查詢工具") as demo:
    gr.Markdown("# 資料庫查詢工具")
    gr.Markdown("這個工具可以幫助您查詢資料庫的知識庫、資料表結構和欄位意義。")
    
    # 獲取所有可用的資料庫
    all_dbs = get_all_databases()
    
    with gr.Tab("知識庫查詢"):
        with gr.Row():
            kb_db = gr.Dropdown(choices=all_dbs, label="選擇資料庫", value=all_dbs[0] if all_dbs else None)
            kb_keyword = gr.Textbox(label="知識關鍵字 (可選)")
        kb_search = gr.Button("查詢知識庫")
        kb_result = gr.DataFrame(label="查詢結果")
        kb_search.click(kb_query, inputs=[kb_db, kb_keyword], outputs=kb_result)
        gr.api(get_all_databases)

    with gr.Tab("資料表查詢"):
        with gr.Row():
            kb_db = gr.Dropdown(choices=all_dbs, label="選擇資料庫", value=all_dbs[0] if all_dbs else None)
        kb_search = gr.Button("查詢資料表")
        kb_result = gr.DataFrame(label="查詢結果")
        kb_search.click(tables_query, inputs=[kb_db], outputs=kb_result)
    
    with gr.Tab("資料表結構查詢"):
        with gr.Row():
            schema_db = gr.Dropdown(choices=all_dbs, label="選擇資料庫", value=all_dbs[0] if all_dbs else None)
            schema_table = gr.Text(label="選擇資料表")
        schema_search = gr.Button("查詢資料表結構")
        schema_result = gr.DataFrame(label="查詢結果")
        
        # 當資料庫選擇變更時，更新資料表下拉選單
        # schema_db.change(update_tables, inputs=schema_db, outputs=schema_table)
        schema_search.click(schema_query, inputs=[schema_db, schema_table], outputs=schema_result)
    
    with gr.Tab("欄位意義查詢"):
        with gr.Row():
            meaning_db = gr.Dropdown(choices=all_dbs, label="選擇資料庫", value=all_dbs[0] if all_dbs else None)
            meaning_table = gr.Text(label="選擇資料表")
        meaning_search = gr.Button("查詢欄位意義")
        meaning_result = gr.DataFrame(label="查詢結果")
        
        # 當資料庫選擇變更時，更新資料表下拉選單
        # meaning_db.change(update_tables, inputs=meaning_db, outputs=meaning_table)
        meaning_search.click(meaning_query, inputs=[meaning_db, meaning_table], outputs=meaning_result)

    with gr.Tab("sqlite SQL 執行"):
        with gr.Row():
            sql_db = gr.Dropdown(choices=all_dbs, label="選擇資料庫", value=all_dbs[0] if all_dbs else None)
            sql_query  = gr.Textbox(label="sqlite SQL")
        sql_exec = gr.Button("執行 SQL")
        sql_result = gr.DataFrame(label="查詢結果")
        sql_exec.click(execute_sqlite, inputs=[sql_db, sql_query], outputs=sql_result)

# 啟動 Gradio 應用程式
if __name__ == "__main__":
    demo.launch(mcp_server=True, server_name="0.0.0.0",allowed_paths=["/"], share=True)