基于 Vanna.AI 和 Amazon Bedrock 构建 Text-to-SQL 方案

CREATE TABLE langchain_pg_embedding (
    uuid UUID PRIMARY KEY,
    collection_id TEXT NOT NULL,
    embedding VECTOR(N), -- N 是嵌入向量的维度，取决于使用的嵌入模型
    document TEXT,
    cmetadata JSONB,
    custom_id TEXT
);

CREATE INDEX ON langchain_pg_embedding
USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 100);

SELECT collection_id, document, cmetadata FROM langchain_pg_embedding;

pip install --upgrade pip setuptools
pip install --upgrade boto3 botocore
pip install boto3 PyMySQL vanna
pip install 'vanna[chromadb,anthropic,postgres,pymysql]'

# 使用默认的embedding model all-MiniLM-L6-v2
import boto3
from vanna.pgvector.pgvector import PG_VectorStore
from vanna.bedrock.bedrock_converse import Bedrock_Converse

# AWS 凭证
AWS_ACCESS_KEY_ID = ' <your-access-key>'
AWS_SECRET_ACCESS_KEY = ' <your-secret_key>'
AWS_REGION = 'us-east-1'

# 创建 Bedrock 客户端
try:
    bedrock_client = boto3.client(
     'bedrock-runtime',
      aws_access_key_id=AWS_ACCESS_KEY_ID,
      aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
      region_name=AWS_REGION
    )
    print(f"✅ Bedrock 客户端创建成功 (区域: {AWS_REGION})")
except Exception as e:
    print(f"❌ Bedrock 客户端创建失败: {str(e)}")

# 数据库连接变量
my_host = 'hostname'
my_dbname = 'pgvector'
my_user = 'username'
my_password = 'password'
my_port = 5432

# 配置 Bedrock 模型用于对话
bedrock_config = {
    "modelId": "us.deepseek.r1-v1:0",
    "allow_llm_to_see_data": True,  # 设置允许LLM查看数据
    "temperature": 0.7,
    "max_tokens": 500
}

# 初始化 MyVanna 实例
try:
    class MyVanna(PG_VectorStore, Bedrock_Converse):
      def __init__(self, config=None):
        if config is None:
            config = {}

        # 构建连接字符串
        config["connection_string"] = f"postgresql://{my_user}:{my_password}@{my_host}:{my_port}/{my_dbname}"

        # 初始化 PG_VectorStore
        PG_VectorStore.__init__(self, config=config)

        # 初始化 Bedrock_Converse
        Bedrock_Converse.__init__(self, client=bedrock_client, config=bedrock_config)
        
# 创建 MyVanna 实例
      vn = MyVanna()
    print(f"✅ MyVanna 实例初始化成功！Vector Database(是: {my_host}:{my_port}/{my_dbname})")
except Exception as e:
    print(f"❌ 连接Vecotr数据库时发生错误: {e}")

#需要被查询的库
try:
    vn.connect_to_postgres(host=my_host, dbname=my_dbname, user=my_user, password=my_password, port=my_port)
    print(f"✅ 成功连接到(数据库: {my_host}:{my_port}/{my_dbname})")
except Exception as e:
    print(f"❌ 连接数据库时发生错误: {e}")

# 打印使用的 embedding model 信息
if hasattr(vn.embedding_function, 'model_id'):
    print(f"使用的 embedding model: {vn.embedding_function.model_id}")
elif hasattr(vn.embedding_function, 'model_name'):
    print(f"使用的 embedding model: {vn.embedding_function.model_name}")
else:
    print(f"使用的 embedding model: 未知 (类型: {type(vn.embedding_function).__name__})")

pip install langchain-aws

#使用amazon.titan-embed-text-v2模型
import boto3
from langchain_aws import BedrockEmbeddings
from vanna.pgvector.pgvector import PG_VectorStore
from vanna.bedrock.bedrock_converse import Bedrock_Converse

# AWS 凭证
AWS_ACCESS_KEY_ID = '<your-access-key>'
AWS_SECRET_ACCESS_KEY = '<your-secret_key>'
AWS_REGION = 'us-east-1'

# 创建 Bedrock 客户端
try:
    bedrock_client = boto3.client(
     'bedrock-runtime',
      aws_access_key_id=AWS_ACCESS_KEY_ID,
      aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
      region_name=AWS_REGION
     )
    print(f"✅ Bedrock 客户端创建成功 (区域: {AWS_REGION})")
except Exception as e:
    print(f"❌ Bedrock 客户端创建失败: {str(e)}")

# 数据库连接变量
my_host = '44.201.54.116'
my_dbname = 'pgvector'
my_user = 'pgvector'
my_password = 'password'
my_port = 54333

# 创建 Bedrock Titan 嵌入模型
embedding_function = BedrockEmbeddings(
    client=bedrock_client,
    model_id="amazon.titan-embed-text-v2:0"  # Titan 嵌入模型
)

# 配置 Bedrock 模型用于对话
bedrock_config = {
    "modelId": "us.deepseek.r1-v1:0",
    "allow_llm_to_see_data": True,  # 设置允许LLM查看数据
    "temperature": 0.7,
    "max_tokens": 500
}

# 初始化 MyVanna 实例
try:
    class MyVanna(PG_VectorStore, Bedrock_Converse):
      def __init__(self, config=None):
        if config is None:
            config = {}

        # 构建连接字符串
        config["connection_string"] = f"postgresql://{my_user}:{my_password}@{my_host}:{my_port}/{my_dbname}"

        # 添加嵌入函数到配置
        config["embedding_function"] = embedding_function

        # 初始化 PG_VectorStore
        PG_VectorStore.__init__(self, config=config)

        # 初始化 Bedrock_Converse
        Bedrock_Converse.__init__(self, client=bedrock_client, config=bedrock_config)

# 创建 MyVanna 实例
      vn = MyVanna()
    print(f"✅ MyVanna 实例初始化成功！Vector Database(是: {my_host}:{my_port}/{my_dbname})")
except Exception as e:
    print(f"❌ 连接Vector数据库时发生错误: {e}")

#需要被查询的库
try:
    vn.connect_to_postgres(host=my_host, dbname=my_dbname, user=my_user, password=my_password, port=my_port)
    print(f"✅ 成功连接到(数据库: {my_host}:{my_port}/{my_dbname})")
except Exception as e:
    print(f"❌ 连接数据库时发生错误: {e}")

# 打印使用的 embedding model 信息
if hasattr(vn.embedding_function, 'model_id'):
    print(f"使用的 embedding model是: {vn.embedding_function.model_id}")
elif hasattr(vn.embedding_function, 'model_name'):
    print(f"使用的 embedding model是: {vn.embedding_function.model_name}")
else:
    print(f"使用的 embedding model是: 未知 (类型: {type(vn.embedding_function).__name__})")

# The information schema query may need some tweaking depending on your database. This is a good starting point.
df_information_schema = vn.run_sql("SELECT * FROM INFORMATION_SCHEMA.COLUMNS")

# This will break up the information schema into bite-sized chunks that can be referenced by the LLM
plan = vn.get_training_plan_generic(df_information_schema)
plan

# If you like the plan, then uncomment this and run it to train
vn.train(plan=plan)

# The following are methods for adding training data. Make sure you modify the examples to match your database.

# DDL statements are powerful because they specify table names, colume names, types, and potentially relationships
vn.train(ddl="""
    CREATE TABLE IF NOT EXISTS my-table (
        id INT PRIMARY KEY,
        name VARCHAR(100),
        age INT
    )
""")

# Sometimes you may want to add documentation about your business terminology or definitions.
vn.train(documentation="有多少个年龄小于30的人?")

# You can also add SQL queries to your training data. This is useful if you have some queries already laying around. You can just copy and paste those from your editor to begin generating new SQL.
vn.train(sql="SELECT count(*) FROM my_table WHERE age < 30")

# At any time you can inspect what training data the package is able to reference
training_data = vn.get_training_data()
training_data

# You can remove training data if there's obsolete/incorrect information. 
vn.remove_training_data(id='7db2674a-a22d-5a3e-8e87-5413cfef9756-sql')

vn.train(sql="SELECT count(*) FROM my_table WHERE age < 30", documentation="有多少个年龄小于30的人?")

vn.ask("针对每个年龄段，统计人数，以10岁为一个年龄段?")

Add of existing embedding ID: d202524f-ce0f-5750-980c-c7cc7c9f69c9-sql
Add of existing embedding ID: 8dde0b98-0a23-50dc-8a90-787057bdb503-sql
Number of requested results 10 is greater than number of elements in index 5, updating n_results = 5
Number of requested results 10 is greater than number of elements in index 1, updating n_results = 1
Add of existing embedding ID: 5c4c38b0-e1fa-5b6d-b782-237117a6eb72-doc
SQL Prompt: [{'role': 'system', 'content': "You are a SQL expert. Please help to generate a SQL query to answer the question. Your response should ONLY be based on the given context and follow the response guidelines and format instructions. \n===Tables \n\n    CREATE TABLE IF NOT EXISTS my-table (\n        id INT PRIMARY KEY,\n        name VARCHAR(100),\n        age INT\n    )\n\n\n\n===Additional Context \n\n有多少个年龄小于30的人?\n\nThe following columns are in the ST_UNITS_OF_MEASURE table in the def database:\n\n|     | TABLE_CATALOG   | TABLE_SCHEMA       | TABLE_NAME          | COLUMN_NAME       | DATA_TYPE   | COLUMN_COMMENT   |\n|----:|:----------------|:-------------------|:--------------------|:------------------|:------------|:-----------------|\n| 741 | def             | information_schema | ST_UNITS_OF_MEASURE | UNIT_NAME         | varchar     |                  |\n| 742 | def             | information_schema | ST_UNITS_OF_MEASURE | UNIT_TYPE         | varchar     |                  |\n| 743 | def             | information_schema | ST_UNITS_OF_MEASURE | CONVERSION_FACTOR | double      |                  |\n| 744 | def             | information_schema | ST_UNITS_OF_MEASURE | DESCRIPTION       | varchar     |                  |\n\nThe following columns are in the x$user_summary_by_file_io table in the def database:\n\n|      | TABLE_CATALOG   | TABLE_SCHEMA   | TABLE_NAME                | COLUMN_NAME   | DATA_TYPE   | COLUMN_COMMENT   |\n|-----:|:----------------|:---------------|:--------------------------|:--------------|:------------|:-----------------|\n| 3754 | def             | sys            | x$user_summary_by_file_io | user          | varchar     |                  |\n| 3755 | def             | sys            | x$user_summary_by_file_io | ios           | decimal     |                  |\n| 3756 | def             | sys            | x$user_summary_by_file_io | io_latency    | decimal     |                  |\n\n===Response Guidelines \n1. If the provided context is sufficient, please generate a valid SQL query without any explanations for the question. \n2. If the provided context is almost sufficient but requires knowledge of a specific string in a particular column, please generate an intermediate SQL query to find the distinct strings in that column. Prepend the query with a comment saying intermediate_sql \n3. If the provided context is insufficient, please explain why it can't be generated. \n4. Please use the most relevant table(s). \n5. If the question has been asked and answered before, please repeat the answer exactly as it was given before. \n6. Ensure that the output SQL is SQL-compliant and executable, and free of syntax errors. \n"}, {'role': 'user', 'content': '针对每个年龄段，统计人数，以10岁为一个年龄段?'}, {'role': 'assistant', 'content': "SELECT \n    FLOOR(age/10)*10 AS age_group,\n    CONCAT(FLOOR(age/10)*10, '-', FLOOR(age/10)*10+9) AS age_range,\n    COUNT(*) AS count\nFROM my_table\nGROUP BY FLOOR(age/10)\nORDER BY age_group\n"}, {'role': 'user', 'content': '列出小于10的人名字和年龄?'}, {'role': 'assistant', 'content': 'SELECT name, age FROM my_table WHERE age < 10'}, {'role': 'user', 'content': '有多少个年龄小于30的人?'}, {'role': 'assistant', 'content': 'SELECT COUNT(*) FROM my_table WHERE age < 30'}, {'role': 'user', 'content': '有多少个年龄小于30的人?'}, {'role': 'assistant', 'content': 'SELECT count(*) FROM my_table WHERE age < 30'}, {'role': 'user', 'content': 'How many individuals are under 30 years old?'}, {'role': 'assistant', 'content': 'SELECT count(*) FROM my_table WHERE age < 30'}, {'role': 'user', 'content': '针对每个年龄段，统计人数，以10岁为一个年龄段?'}]
LLM Response: SELECT 
    FLOOR(age/10)*10 AS age_group,
    CONCAT(FLOOR(age/10)*10, '-', FLOOR(age/10)*10+9) AS age_range,
    COUNT(*) AS count
FROM my_table
GROUP BY FLOOR(age/10)
ORDER BY age_group
SELECT 
    FLOOR(age/10)*10 AS age_group,
    CONCAT(FLOOR(age/10)*10, '-', FLOOR(age/10)*10+9) AS age_range,
    COUNT(*) AS count
FROM my_table
GROUP BY FLOOR(age/10)
ORDER BY age_group
   age_group age_range  count
0          0       0-9     11
1         10     10-19     12
2         20     20-29     12
3         30     30-39     14
4         40     40-49     14
5         50     50-59     12
6         60     60-69     12
7         70     70-79     12
8         80     80-89      1

from vanna.flask import VannaFlaskApp
app = VannaFlaskApp(vn,
                     auth=SimplePassword(users=[{"email": "admin@mail.com", "password": "<your-password>"}]),
                    logo="https://<xxxxxx>.com/16540818023239040",
                    allow_llm_to_see_data=True,
                    summarization=True,
                    title=" XiaoliangGenBI",
                    subtitle="sql 生成器")
app.run(port = 8084, host = '0.0.0.0' , debug = True)