给BI插上AI的翅膀

前言

Superset目前还不擅长关联查询的工作，需要我们在SQL Lab里自己手写各种复杂SQL，本次尝试就2是利用AI技术帮我们写出复杂的SQL。这次尝试比较跨界，使用到了以下技术，

• • Apache Superset

• • DuckDB

• • LLM，主要是Text2SQL应用，当然也可以把之前写的大模型助力情感分析也整合进来。

准备知识

要实现这个功能，我们需要利用到DuckDB的用户自定义函数UDF，把基于大模型的Text2SQL 封装为一个函数，让AI自动帮你写出复杂的SQL。

Text2SQL

为简化工作，这里使用了前几天介绍的vanna.ai，有兴趣的可以看看这几篇文章，

• • 5行代码打造你的大模型数据库聊天机器人

• • 简单代码让你的数据库聊天机器人支持Excel、CSV、JSON、Parquet甚至数据湖

• • 不同 LLMs 和上下文策略的测试，以最大化 SQL 生成准确性

• • 数据库聊天机器人之本地化部署

DuckDB UDF

如果还不太熟悉DuckDB UDF的，可以看看这几篇文章，

• • 使用 DuckDB + Ibis 进行 RAG

• • 探索DuckDB：使用标量 Python UDF 快速扩展 DuckDB 的功能

• • 一句 SQL生成海量超真实模拟数据

duckdb_engine

当然 duckdb_engine
 默认还不能加载UDF，需要做些修改，这些代码我放到了我的 repo^[1] 里了，

演示用代码

为了简化工作，这里使用vanna.ai， 如果愿意可以结合我之前的文章，使用本地LLM（私有化部署）

udf.py
内容如下，

import pyarrow.compute as pc
import duckdb
import pyarrow as pa

from vanna.remote import VannaDefault

api_key = os.getenv('VANNA_API_KEY') # Your API key from https://vanna.ai/account/profile 
vanna_model_name = 'tpc-duck' # Your model name from https://vanna.ai/account/profile
                    
vn = VannaDefault(model=vanna_model_name, api_key=api_key)

def connect_to_duckdb1(conn):
    conn = conn

    def run_sql_duckdb(sql: str):
        return conn.query(sql).to_df()

    # global run_sql
    vn.run_sql = run_sql_duckdb
vn.connect_to_duckdb=connect_to_duckdb1

def exist_function(conn,name):
    function_check = f"""SELECT DISTINCT  function_name
                        FROM duckdb_functions()
                        WHERE lower(function_type) = 'scalar'
                        AND lower(function_name) in ('{name.lower()}')
                        ORDER BY function_name;"""

    function_check_output = conn.query(function_check)
    return function_check_output

def ai(x):
    sql = vn.ask(x,False,False,False)
    return sql[0]

def create_functions(conn):
    if not exist_function(conn,'ai'):
        conn.create_function('ai', ai, ['VARCHAR'], 'VARCHAR', type='native')

结论

当前这个为简化版本，对Superset本身没有做什么修改。从AI辅助对角度可以做的工作还有很多。另外基于这样的方法实现，我们不仅仅可以把AI整合进Superset，还可以把各种机器学习模型也整合进来，甚至ETL工作等。

引用链接

[1]
 repo: http://github.com/alitrack/duckdb_engine