Initial commit
This commit is contained in:
@@ -0,0 +1,17 @@
|
||||
"""
|
||||
简化的Text2SQL框架
|
||||
基于RAGFlow方案实现的Text2SQL框架
|
||||
"""
|
||||
|
||||
__version__ = "1.0.0"
|
||||
__author__ = "RAG Team"
|
||||
|
||||
from .knowledge_base import SimpleKnowledgeBase
|
||||
from .sql_generator import SimpleSQLGenerator
|
||||
from .text2sql_agent import SimpleText2SQLAgent
|
||||
|
||||
__all__ = [
|
||||
"SimpleKnowledgeBase",
|
||||
"SimpleSQLGenerator",
|
||||
"SimpleText2SQLAgent"
|
||||
]
|
||||
@@ -0,0 +1,57 @@
|
||||
[
|
||||
{
|
||||
"table_name": "users",
|
||||
"table_description": "用户信息表，存储注册用户的基本信息",
|
||||
"columns": [
|
||||
{"name": "id", "description": "用户唯一标识符，主键", "type": "INT"},
|
||||
{"name": "name", "description": "用户姓名，不能为空", "type": "VARCHAR(100)"},
|
||||
{"name": "email", "description": "用户邮箱地址，必须唯一", "type": "VARCHAR(150)"},
|
||||
{"name": "age", "description": "用户年龄", "type": "INT"},
|
||||
{"name": "created_at", "description": "用户注册时间", "type": "TIMESTAMP"}
|
||||
]
|
||||
},
|
||||
{
|
||||
"table_name": "orders",
|
||||
"table_description": "订单表，记录用户的购买订单信息",
|
||||
"columns": [
|
||||
{"name": "id", "description": "订单唯一标识符，主键", "type": "INT"},
|
||||
{"name": "user_id", "description": "下单用户的ID，外键关联users表", "type": "INT"},
|
||||
{"name": "product_name", "description": "购买的产品名称", "type": "VARCHAR(200)"},
|
||||
{"name": "quantity", "description": "购买数量", "type": "INT"},
|
||||
{"name": "price", "description": "订单总价格", "type": "DECIMAL(10,2)"},
|
||||
{"name": "order_date", "description": "下单时间", "type": "TIMESTAMP"}
|
||||
]
|
||||
},
|
||||
{
|
||||
"table_name": "products",
|
||||
"table_description": "产品表，存储商城中所有产品的信息",
|
||||
"columns": [
|
||||
{"name": "id", "description": "产品唯一标识符，主键", "type": "INT"},
|
||||
{"name": "name", "description": "产品名称", "type": "VARCHAR(200)"},
|
||||
{"name": "category", "description": "产品分类", "type": "VARCHAR(100)"},
|
||||
{"name": "price", "description": "产品单价", "type": "DECIMAL(10,2)"},
|
||||
{"name": "stock", "description": "库存数量", "type": "INT"},
|
||||
{"name": "description", "description": "产品详细描述", "type": "TEXT"}
|
||||
]
|
||||
},
|
||||
{
|
||||
"table_name": "categories",
|
||||
"table_description": "产品分类表，定义产品的分类信息",
|
||||
"columns": [
|
||||
{"name": "id", "description": "分类唯一标识符，主键", "type": "INT"},
|
||||
{"name": "name", "description": "分类名称，必须唯一", "type": "VARCHAR(100)"},
|
||||
{"name": "description", "description": "分类描述", "type": "TEXT"}
|
||||
]
|
||||
},
|
||||
{
|
||||
"table_name": "order_items",
|
||||
"table_description": "订单明细表，存储订单中包含的具体商品信息",
|
||||
"columns": [
|
||||
{"name": "id", "description": "订单明细唯一标识符，主键", "type": "INT"},
|
||||
{"name": "order_id", "description": "关联的订单ID，外键", "type": "INT"},
|
||||
{"name": "product_id", "description": "关联的产品ID，外键", "type": "INT"},
|
||||
{"name": "quantity", "description": "该商品在订单中的数量", "type": "INT"},
|
||||
{"name": "unit_price", "description": "该商品的单价（下单时的价格）", "type": "DECIMAL(10,2)"}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,27 @@
|
||||
[
|
||||
{
|
||||
"table_name": "users",
|
||||
"ddl_statement": "CREATE TABLE users (id INT PRIMARY KEY AUTO_INCREMENT, name VARCHAR(100) NOT NULL, email VARCHAR(150) UNIQUE, age INT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)",
|
||||
"description": "用户信息表，存储用户基本信息"
|
||||
},
|
||||
{
|
||||
"table_name": "orders",
|
||||
"ddl_statement": "CREATE TABLE orders (id INT PRIMARY KEY AUTO_INCREMENT, user_id INT, product_name VARCHAR(200), quantity INT, price DECIMAL(10,2), order_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY (user_id) REFERENCES users(id))",
|
||||
"description": "订单表，存储用户订单信息"
|
||||
},
|
||||
{
|
||||
"table_name": "products",
|
||||
"ddl_statement": "CREATE TABLE products (id INT PRIMARY KEY AUTO_INCREMENT, name VARCHAR(200) NOT NULL, category VARCHAR(100), price DECIMAL(10,2), stock INT DEFAULT 0, description TEXT)",
|
||||
"description": "产品表，存储产品基本信息"
|
||||
},
|
||||
{
|
||||
"table_name": "categories",
|
||||
"ddl_statement": "CREATE TABLE categories (id INT PRIMARY KEY AUTO_INCREMENT, name VARCHAR(100) NOT NULL UNIQUE, description TEXT)",
|
||||
"description": "产品分类表"
|
||||
},
|
||||
{
|
||||
"table_name": "order_items",
|
||||
"ddl_statement": "CREATE TABLE order_items (id INT PRIMARY KEY AUTO_INCREMENT, order_id INT, product_id INT, quantity INT, unit_price DECIMAL(10,2), FOREIGN KEY (order_id) REFERENCES orders(id), FOREIGN KEY (product_id) REFERENCES products(id))",
|
||||
"description": "订单明细表，存储订单中的具体商品信息"
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,62 @@
|
||||
[
|
||||
{
|
||||
"question": "查询所有用户的姓名和邮箱",
|
||||
"sql": "SELECT name, email FROM users",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查找年龄大于25岁的用户",
|
||||
"sql": "SELECT * FROM users WHERE age > 25",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查询每个用户的订单数量",
|
||||
"sql": "SELECT u.name, COUNT(o.id) as order_count FROM users u LEFT JOIN orders o ON u.id = o.user_id GROUP BY u.id, u.name",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查找最近7天的订单",
|
||||
"sql": "SELECT * FROM orders WHERE order_date >= DATE_SUB(NOW(), INTERVAL 7 DAY)",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查询总销售额最高的前5个产品",
|
||||
"sql": "SELECT p.name, SUM(oi.quantity * oi.unit_price) as total_sales FROM products p JOIN order_items oi ON p.id = oi.product_id GROUP BY p.id, p.name ORDER BY total_sales DESC LIMIT 5",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查询某个用户的所有订单",
|
||||
"sql": "SELECT o.*, u.name as user_name FROM orders o JOIN users u ON o.user_id = u.id WHERE u.name = '张三'",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查询价格在100到500之间的产品",
|
||||
"sql": "SELECT * FROM products WHERE price BETWEEN 100 AND 500",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查询库存少于10的产品",
|
||||
"sql": "SELECT * FROM products WHERE stock < 10",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查询每个分类的产品数量",
|
||||
"sql": "SELECT category, COUNT(*) as product_count FROM products GROUP BY category",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查询订单总金额大于1000的订单",
|
||||
"sql": "SELECT o.*, SUM(oi.quantity * oi.unit_price) as total_amount FROM orders o JOIN order_items oi ON o.id = oi.order_id GROUP BY o.id HAVING total_amount > 1000",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查询没有下过订单的用户",
|
||||
"sql": "SELECT u.* FROM users u LEFT JOIN orders o ON u.id = o.user_id WHERE o.id IS NULL",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查询平均订单金额",
|
||||
"sql": "SELECT AVG(total_amount) as avg_order_amount FROM (SELECT o.id, SUM(oi.quantity * oi.unit_price) as total_amount FROM orders o JOIN order_items oi ON o.id = oi.order_id GROUP BY o.id) as order_totals",
|
||||
"database": "ecommerce"
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,184 @@
|
||||
import json
|
||||
import os
|
||||
from typing import List, Dict, Any
|
||||
from pymilvus import MilvusClient, FieldSchema, CollectionSchema, DataType
|
||||
from pymilvus.model.hybrid import BGEM3EmbeddingFunction
|
||||
|
||||
|
||||
class SimpleKnowledgeBase:
|
||||
"""知识库"""
|
||||
|
||||
def __init__(self, milvus_uri: str = "http://localhost:19530"):
|
||||
self.milvus_uri = milvus_uri
|
||||
self.client = MilvusClient(uri=milvus_uri)
|
||||
self.embedding_function = BGEM3EmbeddingFunction(use_fp16=False, device="cpu")
|
||||
self.collection_name = "text2sql_kb"
|
||||
self._setup_collection()
|
||||
|
||||
def _setup_collection(self):
|
||||
"""设置集合"""
|
||||
if self.client.has_collection(self.collection_name):
|
||||
self.client.drop_collection(self.collection_name)
|
||||
|
||||
# 定义字段
|
||||
fields = [
|
||||
FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=True, max_length=100),
|
||||
FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=4096),
|
||||
FieldSchema(name="type", dtype=DataType.VARCHAR, max_length=32), # ddl, qsql, description
|
||||
FieldSchema(name="dense_vector", dtype=DataType.FLOAT_VECTOR, dim=self.embedding_function.dim["dense"])
|
||||
]
|
||||
|
||||
schema = CollectionSchema(fields, description="Text2SQL知识库")
|
||||
|
||||
# 创建集合
|
||||
self.client.create_collection(
|
||||
collection_name=self.collection_name,
|
||||
schema=schema,
|
||||
consistency_level="Strong"
|
||||
)
|
||||
|
||||
# 创建索引
|
||||
index_params = self.client.prepare_index_params()
|
||||
index_params.add_index(
|
||||
field_name="dense_vector",
|
||||
index_type="AUTOINDEX",
|
||||
metric_type="IP"
|
||||
)
|
||||
|
||||
self.client.create_index(
|
||||
collection_name=self.collection_name,
|
||||
index_params=index_params
|
||||
)
|
||||
|
||||
def load_data(self):
|
||||
"""加载所有知识库数据"""
|
||||
data_dir = os.path.join(os.path.dirname(__file__), "data")
|
||||
|
||||
# 加载DDL数据
|
||||
ddl_path = os.path.join(data_dir, "ddl_examples.json")
|
||||
if os.path.exists(ddl_path):
|
||||
with open(ddl_path, 'r', encoding='utf-8') as f:
|
||||
ddl_data = json.load(f)
|
||||
self._add_ddl_data(ddl_data)
|
||||
|
||||
# 加载Q->SQL数据
|
||||
qsql_path = os.path.join(data_dir, "qsql_examples.json")
|
||||
if os.path.exists(qsql_path):
|
||||
with open(qsql_path, 'r', encoding='utf-8') as f:
|
||||
qsql_data = json.load(f)
|
||||
self._add_qsql_data(qsql_data)
|
||||
|
||||
# 加载描述数据
|
||||
desc_path = os.path.join(data_dir, "db_descriptions.json")
|
||||
if os.path.exists(desc_path):
|
||||
with open(desc_path, 'r', encoding='utf-8') as f:
|
||||
desc_data = json.load(f)
|
||||
self._add_description_data(desc_data)
|
||||
|
||||
# 加载集合到内存
|
||||
self.client.load_collection(collection_name=self.collection_name)
|
||||
print("知识库数据加载完成")
|
||||
|
||||
def _add_ddl_data(self, data: List[Dict]):
|
||||
"""添加DDL数据"""
|
||||
contents = []
|
||||
types = []
|
||||
|
||||
for item in data:
|
||||
content = f"表名: {item.get('table_name', '')}\n"
|
||||
content += f"DDL: {item.get('ddl_statement', '')}\n"
|
||||
content += f"描述: {item.get('description', '')}"
|
||||
|
||||
contents.append(content)
|
||||
types.append("ddl")
|
||||
|
||||
self._insert_data(contents, types)
|
||||
|
||||
def _add_qsql_data(self, data: List[Dict]):
|
||||
"""添加Q->SQL数据"""
|
||||
contents = []
|
||||
types = []
|
||||
|
||||
for item in data:
|
||||
content = f"问题: {item.get('question', '')}\n"
|
||||
content += f"SQL: {item.get('sql', '')}"
|
||||
|
||||
contents.append(content)
|
||||
types.append("qsql")
|
||||
|
||||
self._insert_data(contents, types)
|
||||
|
||||
def _add_description_data(self, data: List[Dict]):
|
||||
"""添加描述数据"""
|
||||
contents = []
|
||||
types = []
|
||||
|
||||
for item in data:
|
||||
content = f"表名: {item.get('table_name', '')}\n"
|
||||
content += f"表描述: {item.get('table_description', '')}\n"
|
||||
|
||||
columns = item.get('columns', [])
|
||||
if columns:
|
||||
content += "字段信息:\n"
|
||||
for col in columns:
|
||||
content += f" - {col.get('name', '')}: {col.get('description', '')} ({col.get('type', '')})\n"
|
||||
|
||||
contents.append(content)
|
||||
types.append("description")
|
||||
|
||||
self._insert_data(contents, types)
|
||||
|
||||
def _insert_data(self, contents: List[str], types: List[str]):
|
||||
"""插入数据"""
|
||||
if not contents:
|
||||
return
|
||||
|
||||
# 生成嵌入
|
||||
embeddings = self.embedding_function(contents)
|
||||
|
||||
# 构建插入数据，每一行是一个字典
|
||||
data_to_insert = []
|
||||
for i in range(len(contents)):
|
||||
data_to_insert.append({
|
||||
"content": contents[i],
|
||||
"type": types[i],
|
||||
"dense_vector": embeddings["dense"][i]
|
||||
})
|
||||
|
||||
# 插入数据
|
||||
result = self.client.insert(
|
||||
collection_name=self.collection_name,
|
||||
data=data_to_insert
|
||||
)
|
||||
|
||||
def search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
|
||||
"""搜索相关内容"""
|
||||
self.client.load_collection(collection_name=self.collection_name)
|
||||
|
||||
query_embeddings = self.embedding_function([query])
|
||||
|
||||
search_results = self.client.search(
|
||||
collection_name=self.collection_name,
|
||||
data=query_embeddings["dense"],
|
||||
anns_field="dense_vector",
|
||||
search_params={"metric_type": "IP"},
|
||||
limit=top_k,
|
||||
output_fields=["content", "type"]
|
||||
)
|
||||
|
||||
results = []
|
||||
for hit in search_results[0]:
|
||||
results.append({
|
||||
"content": hit["entity"]["content"],
|
||||
"type": hit["entity"]["type"],
|
||||
"score": hit["distance"]
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
def cleanup(self):
|
||||
"""清理资源"""
|
||||
try:
|
||||
self.client.drop_collection(self.collection_name)
|
||||
except:
|
||||
pass
|
||||
@@ -0,0 +1,113 @@
|
||||
import os
|
||||
from typing import List, Dict, Any
|
||||
from langchain_deepseek import ChatDeepSeek
|
||||
from langchain.schema import HumanMessage, SystemMessage
|
||||
|
||||
|
||||
class SimpleSQLGenerator:
|
||||
"""简化的SQL生成器"""
|
||||
|
||||
def __init__(self, api_key: str = None):
|
||||
self.llm = ChatDeepSeek(
|
||||
model="deepseek-chat",
|
||||
temperature=0,
|
||||
api_key=api_key or os.getenv("DEEPSEEK_API_KEY")
|
||||
)
|
||||
|
||||
def generate_sql(self, user_query: str, knowledge_results: List[Dict[str, Any]]) -> str:
|
||||
"""生成SQL语句"""
|
||||
# 构建上下文
|
||||
context = self._build_context(knowledge_results)
|
||||
|
||||
# 构建提示
|
||||
prompt = f"""你是一个SQL专家。请根据以下信息将用户问题转换为SQL查询语句。
|
||||
|
||||
数据库信息：
|
||||
{context}
|
||||
|
||||
用户问题：{user_query}
|
||||
|
||||
要求：
|
||||
1. 只返回SQL语句，不要包含任何解释
|
||||
2. 确保SQL语法正确
|
||||
3. 使用上下文中提供的表名和字段名
|
||||
4. 如果需要JOIN，请根据表结构进行合理关联
|
||||
|
||||
SQL语句："""
|
||||
|
||||
messages = [HumanMessage(content=prompt)]
|
||||
response = self.llm.invoke(messages)
|
||||
|
||||
# 清理SQL语句
|
||||
sql = response.content.strip()
|
||||
if sql.startswith("```sql"):
|
||||
sql = sql[6:]
|
||||
if sql.startswith("```"):
|
||||
sql = sql[3:]
|
||||
if sql.endswith("```"):
|
||||
sql = sql[:-3]
|
||||
|
||||
return sql.strip()
|
||||
|
||||
def fix_sql(self, original_sql: str, error_message: str, knowledge_results: List[Dict[str, Any]]) -> str:
|
||||
"""修复SQL语句"""
|
||||
context = self._build_context(knowledge_results)
|
||||
|
||||
prompt = f"""请修复以下SQL语句的错误。
|
||||
|
||||
数据库信息：
|
||||
{context}
|
||||
|
||||
原始SQL：
|
||||
{original_sql}
|
||||
|
||||
错误信息：
|
||||
{error_message}
|
||||
|
||||
请返回修复后的SQL语句（只返回SQL，不要解释）："""
|
||||
|
||||
messages = [HumanMessage(content=prompt)]
|
||||
response = self.llm.invoke(messages)
|
||||
|
||||
# 清理SQL语句
|
||||
fixed_sql = response.content.strip()
|
||||
if fixed_sql.startswith("```sql"):
|
||||
fixed_sql = fixed_sql[6:]
|
||||
if fixed_sql.startswith("```"):
|
||||
fixed_sql = fixed_sql[3:]
|
||||
if fixed_sql.endswith("```"):
|
||||
fixed_sql = fixed_sql[:-3]
|
||||
|
||||
return fixed_sql.strip()
|
||||
|
||||
def _build_context(self, knowledge_results: List[Dict[str, Any]]) -> str:
|
||||
"""构建上下文信息"""
|
||||
context = ""
|
||||
|
||||
# 按类型分组
|
||||
ddl_info = []
|
||||
qsql_examples = []
|
||||
descriptions = []
|
||||
|
||||
for result in knowledge_results:
|
||||
if result["type"] == "ddl":
|
||||
ddl_info.append(result["content"])
|
||||
elif result["type"] == "qsql":
|
||||
qsql_examples.append(result["content"])
|
||||
elif result["type"] == "description":
|
||||
descriptions.append(result["content"])
|
||||
|
||||
# 构建上下文
|
||||
if ddl_info:
|
||||
context += "=== 表结构信息 ===\n"
|
||||
context += "\n".join(ddl_info) + "\n\n"
|
||||
|
||||
if descriptions:
|
||||
context += "=== 表和字段描述 ===\n"
|
||||
context += "\n".join(descriptions) + "\n\n"
|
||||
|
||||
if qsql_examples:
|
||||
context += "=== 查询示例 ===\n"
|
||||
context += "\n".join(qsql_examples) + "\n\n"
|
||||
|
||||
return context
|
||||
@@ -0,0 +1,213 @@
|
||||
import sqlite3
|
||||
import os
|
||||
from typing import Dict, Any, List, Tuple
|
||||
from .knowledge_base import SimpleKnowledgeBase
|
||||
from .sql_generator import SimpleSQLGenerator
|
||||
|
||||
|
||||
class SimpleText2SQLAgent:
|
||||
"""Text2SQL代理"""
|
||||
|
||||
def __init__(self, milvus_uri: str = "http://localhost:19530", api_key: str = None):
|
||||
"""初始化代理"""
|
||||
self.knowledge_base = SimpleKnowledgeBase(milvus_uri)
|
||||
self.sql_generator = SimpleSQLGenerator(api_key)
|
||||
self.db_path = None
|
||||
self.connection = None
|
||||
|
||||
# 配置参数
|
||||
self.max_retry_count = 3
|
||||
self.top_k_retrieval = 5
|
||||
self.max_result_rows = 100
|
||||
|
||||
def connect_database(self, db_path: str) -> bool:
|
||||
"""连接SQLite数据库"""
|
||||
try:
|
||||
self.db_path = db_path
|
||||
self.connection = sqlite3.connect(db_path)
|
||||
print(f"成功连接到数据库: {db_path}")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"数据库连接失败: {str(e)}")
|
||||
return False
|
||||
|
||||
def load_knowledge_base(self):
|
||||
"""加载知识库"""
|
||||
self.knowledge_base.load_data()
|
||||
|
||||
def query(self, user_question: str) -> Dict[str, Any]:
|
||||
"""执行Text2SQL查询"""
|
||||
if not self.connection:
|
||||
return {
|
||||
"success": False,
|
||||
"error": "数据库未连接",
|
||||
"sql": None,
|
||||
"results": None
|
||||
}
|
||||
|
||||
print(f"\n=== 处理查询: {user_question} ===")
|
||||
|
||||
# 1. 从知识库检索
|
||||
print("检索知识库...")
|
||||
knowledge_results = self.knowledge_base.search(user_question, self.top_k_retrieval)
|
||||
print(f"检索到 {len(knowledge_results)} 条相关信息")
|
||||
|
||||
# 2. 生成SQL
|
||||
print("生成SQL...")
|
||||
sql = self.sql_generator.generate_sql(user_question, knowledge_results)
|
||||
print(f"生成的SQL: {sql}")
|
||||
|
||||
# 3. 执行SQL（带重试）
|
||||
retry_count = 0
|
||||
while retry_count < self.max_retry_count:
|
||||
print(f"执行SQL (尝试 {retry_count + 1}/{self.max_retry_count})...")
|
||||
|
||||
success, result = self._execute_sql(sql)
|
||||
|
||||
if success:
|
||||
print("SQL执行成功!")
|
||||
return {
|
||||
"success": True,
|
||||
"error": None,
|
||||
"sql": sql,
|
||||
"results": result,
|
||||
"retry_count": retry_count
|
||||
}
|
||||
else:
|
||||
print(f"SQL执行失败: {result}")
|
||||
|
||||
if retry_count < self.max_retry_count - 1:
|
||||
print("尝试修复SQL...")
|
||||
sql = self.sql_generator.fix_sql(sql, result, knowledge_results)
|
||||
print(f"修复后的SQL: {sql}")
|
||||
|
||||
retry_count += 1
|
||||
|
||||
return {
|
||||
"success": False,
|
||||
"error": f"超过最大重试次数 ({self.max_retry_count})",
|
||||
"sql": sql,
|
||||
"results": None,
|
||||
"retry_count": retry_count
|
||||
}
|
||||
|
||||
def _execute_sql(self, sql: str) -> Tuple[bool, Any]:
|
||||
"""执行SQL语句"""
|
||||
try:
|
||||
cursor = self.connection.cursor()
|
||||
|
||||
# 添加LIMIT限制
|
||||
if sql.strip().upper().startswith('SELECT') and 'LIMIT' not in sql.upper():
|
||||
sql = f"{sql.rstrip(';')} LIMIT {self.max_result_rows}"
|
||||
|
||||
cursor.execute(sql)
|
||||
|
||||
if sql.strip().upper().startswith('SELECT'):
|
||||
# 查询语句
|
||||
columns = [desc[0] for desc in cursor.description]
|
||||
rows = cursor.fetchall()
|
||||
|
||||
results = []
|
||||
for row in rows:
|
||||
result_row = {}
|
||||
for i, value in enumerate(row):
|
||||
result_row[columns[i]] = value
|
||||
results.append(result_row)
|
||||
|
||||
cursor.close()
|
||||
return True, {
|
||||
"columns": columns,
|
||||
"rows": results,
|
||||
"count": len(results)
|
||||
}
|
||||
else:
|
||||
# 非查询语句
|
||||
self.connection.commit()
|
||||
cursor.close()
|
||||
return True, "SQL执行成功"
|
||||
|
||||
except Exception as e:
|
||||
return False, str(e)
|
||||
|
||||
def add_example(self, question: str, sql: str):
|
||||
"""添加新的Q->SQL示例"""
|
||||
# 简化版本：直接保存到文件
|
||||
data_dir = os.path.join(os.path.dirname(__file__), "data")
|
||||
qsql_path = os.path.join(data_dir, "qsql_examples.json")
|
||||
|
||||
try:
|
||||
import json
|
||||
|
||||
# 读取现有数据
|
||||
if os.path.exists(qsql_path):
|
||||
with open(qsql_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
else:
|
||||
data = []
|
||||
|
||||
# 添加新示例
|
||||
data.append({
|
||||
"question": question,
|
||||
"sql": sql,
|
||||
"database": "sqlite"
|
||||
})
|
||||
|
||||
# 保存
|
||||
with open(qsql_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||||
|
||||
print(f"已添加新示例: {question}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"添加示例失败: {str(e)}")
|
||||
|
||||
def get_table_info(self) -> List[Dict[str, Any]]:
|
||||
"""获取数据库表信息"""
|
||||
if not self.connection:
|
||||
return []
|
||||
|
||||
try:
|
||||
cursor = self.connection.cursor()
|
||||
|
||||
# 获取所有表名
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
|
||||
tables = cursor.fetchall()
|
||||
|
||||
table_info = []
|
||||
for table in tables:
|
||||
table_name = table[0]
|
||||
|
||||
# 获取表结构
|
||||
cursor.execute(f"PRAGMA table_info({table_name})")
|
||||
columns = cursor.fetchall()
|
||||
|
||||
table_info.append({
|
||||
"table_name": table_name,
|
||||
"columns": [
|
||||
{
|
||||
"name": col[1],
|
||||
"type": col[2],
|
||||
"nullable": not col[3],
|
||||
"default": col[4],
|
||||
"primary_key": bool(col[5])
|
||||
}
|
||||
for col in columns
|
||||
]
|
||||
})
|
||||
|
||||
cursor.close()
|
||||
return table_info
|
||||
|
||||
except Exception as e:
|
||||
print(f"获取表信息失败: {str(e)}")
|
||||
return []
|
||||
|
||||
def cleanup(self):
|
||||
"""清理资源"""
|
||||
if self.connection:
|
||||
self.connection.close()
|
||||
self.connection = None
|
||||
print("数据库连接已关闭")
|
||||
|
||||
self.knowledge_base.cleanup()
|
||||
print("知识库已清理")
|
||||
Reference in New Issue
Block a user