Initial commit
This commit is contained in:
@@ -0,0 +1,377 @@
|
||||
import os
|
||||
import json
|
||||
import sqlite3
|
||||
import numpy as np
|
||||
from typing import List, Dict, Any
|
||||
from sentence_transformers import SentenceTransformer
|
||||
from sklearn.feature_extraction.text import TfidfVectorizer
|
||||
from sklearn.metrics.pairwise import cosine_similarity
|
||||
from pymilvus import connections, MilvusClient, FieldSchema, CollectionSchema, DataType, Collection
|
||||
|
||||
|
||||
class BGESmallEmbeddingFunction:
|
||||
"""BGE-Small中文嵌入函数，用于Text2SQL知识库向量化"""
|
||||
|
||||
def __init__(self, model_name="BAAI/bge-small-zh-v1.5", device="cpu"):
|
||||
self.model_name = model_name
|
||||
self.device = device
|
||||
self.model = SentenceTransformer(model_name, device=device)
|
||||
self.dense_dim = self.model.get_sentence_embedding_dimension()
|
||||
|
||||
def encode_text(self, texts):
|
||||
"""编码文本为密集向量"""
|
||||
if isinstance(texts, str):
|
||||
texts = [texts]
|
||||
|
||||
embeddings = self.model.encode(
|
||||
texts,
|
||||
normalize_embeddings=True,
|
||||
batch_size=16,
|
||||
convert_to_numpy=True
|
||||
)
|
||||
|
||||
return embeddings
|
||||
|
||||
@property
|
||||
def dim(self):
|
||||
"""返回向量维度"""
|
||||
return self.dense_dim
|
||||
|
||||
|
||||
class SimpleKnowledgeBase:
|
||||
"""简化的知识库，使用BGE-Small进行向量检索"""
|
||||
|
||||
def __init__(self, milvus_uri: str = "http://localhost:19530"):
|
||||
self.milvus_uri = milvus_uri
|
||||
self.collection_name = "text2sql_knowledge_base"
|
||||
self.milvus_client = None
|
||||
self.collection = None
|
||||
|
||||
self.embedding_function = BGESmallEmbeddingFunction(
|
||||
model_name="BAAI/bge-small-zh-v1.5",
|
||||
device="cpu"
|
||||
)
|
||||
|
||||
self.sql_examples = []
|
||||
self.table_schemas = []
|
||||
self.data_loaded = False
|
||||
|
||||
def connect_milvus(self):
|
||||
"""连接Milvus数据库"""
|
||||
connections.connect(uri=self.milvus_uri)
|
||||
self.milvus_client = MilvusClient(uri=self.milvus_uri)
|
||||
return True
|
||||
|
||||
def create_collection(self):
|
||||
"""创建Milvus集合"""
|
||||
if not self.milvus_client:
|
||||
self.connect_milvus()
|
||||
|
||||
if self.milvus_client.has_collection(self.collection_name):
|
||||
self.milvus_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_type", dtype=DataType.VARCHAR, max_length=50),
|
||||
FieldSchema(name="question", dtype=DataType.VARCHAR, max_length=1000),
|
||||
FieldSchema(name="sql", dtype=DataType.VARCHAR, max_length=2000),
|
||||
FieldSchema(name="description", dtype=DataType.VARCHAR, max_length=1000),
|
||||
FieldSchema(name="table_name", dtype=DataType.VARCHAR, max_length=100),
|
||||
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=self.embedding_function.dim)
|
||||
]
|
||||
|
||||
schema = CollectionSchema(fields, description="Text2SQL知识库")
|
||||
self.collection = Collection(name=self.collection_name, schema=schema, consistency_level="Strong")
|
||||
|
||||
index_params = {"index_type": "AUTOINDEX", "metric_type": "IP", "params": {}}
|
||||
self.collection.create_index("embedding", index_params)
|
||||
|
||||
return True
|
||||
|
||||
def load_data(self):
|
||||
"""加载知识库数据"""
|
||||
data_dir = os.path.join(os.path.dirname(__file__), "data")
|
||||
|
||||
self.load_sql_examples(data_dir)
|
||||
self.load_table_schemas(data_dir)
|
||||
self.vectorize_and_store()
|
||||
|
||||
self.data_loaded = True
|
||||
|
||||
def load_sql_examples(self, data_dir: str):
|
||||
"""加载SQL示例"""
|
||||
sql_examples_path = os.path.join(data_dir, "qsql_examples.json")
|
||||
|
||||
default_examples = [
|
||||
{"question": "查询所有用户信息", "sql": "SELECT * FROM users", "description": "获取用户记录", "database": "sqlite"},
|
||||
{"question": "年龄大于30的用户", "sql": "SELECT * FROM users WHERE age > 30", "description": "年龄筛选", "database": "sqlite"},
|
||||
{"question": "统计用户总数", "sql": "SELECT COUNT(*) as user_count FROM users", "description": "用户计数", "database": "sqlite"},
|
||||
{"question": "查询库存不足的产品", "sql": "SELECT * FROM products WHERE stock < 50", "description": "库存筛选", "database": "sqlite"},
|
||||
{"question": "查询用户订单信息", "sql": "SELECT u.name, p.name, o.quantity FROM orders o JOIN users u ON o.user_id = u.id JOIN products p ON o.product_id = p.id", "description": "订单详情", "database": "sqlite"},
|
||||
{"question": "按城市统计用户", "sql": "SELECT city, COUNT(*) as count FROM users GROUP BY city", "description": "城市分组", "database": "sqlite"}
|
||||
]
|
||||
|
||||
if os.path.exists(sql_examples_path):
|
||||
with open(sql_examples_path, 'r', encoding='utf-8') as f:
|
||||
self.sql_examples = json.load(f)
|
||||
else:
|
||||
self.sql_examples = default_examples
|
||||
os.makedirs(data_dir, exist_ok=True)
|
||||
with open(sql_examples_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(self.sql_examples, f, ensure_ascii=False, indent=2)
|
||||
|
||||
def load_table_schemas(self, data_dir: str):
|
||||
"""加载表结构信息"""
|
||||
schema_path = os.path.join(data_dir, "table_schemas.json")
|
||||
|
||||
default_schemas = [
|
||||
{
|
||||
"table_name": "users",
|
||||
"description": "用户信息表",
|
||||
"columns": [
|
||||
{"name": "id", "type": "INTEGER", "description": "用户ID"},
|
||||
{"name": "name", "type": "VARCHAR", "description": "用户姓名"},
|
||||
{"name": "age", "type": "INTEGER", "description": "用户年龄"},
|
||||
{"name": "email", "type": "VARCHAR", "description": "邮箱地址"},
|
||||
{"name": "city", "type": "VARCHAR", "description": "所在城市"},
|
||||
{"name": "created_at", "type": "DATETIME", "description": "创建时间"}
|
||||
]
|
||||
},
|
||||
{
|
||||
"table_name": "products",
|
||||
"description": "产品信息表",
|
||||
"columns": [
|
||||
{"name": "id", "type": "INTEGER", "description": "产品ID"},
|
||||
{"name": "product_name", "type": "VARCHAR", "description": "产品名称"},
|
||||
{"name": "category", "type": "VARCHAR", "description": "产品类别"},
|
||||
{"name": "price", "type": "DECIMAL", "description": "产品价格"},
|
||||
{"name": "stock", "type": "INTEGER", "description": "库存数量"},
|
||||
{"name": "description", "type": "TEXT", "description": "产品描述"}
|
||||
]
|
||||
},
|
||||
{
|
||||
"table_name": "orders",
|
||||
"description": "订单信息表",
|
||||
"columns": [
|
||||
{"name": "id", "type": "INTEGER", "description": "订单ID"},
|
||||
{"name": "user_id", "type": "INTEGER", "description": "用户ID"},
|
||||
{"name": "product_id", "type": "INTEGER", "description": "产品ID"},
|
||||
{"name": "quantity", "type": "INTEGER", "description": "购买数量"},
|
||||
{"name": "total_price", "type": "DECIMAL", "description": "总价格"},
|
||||
{"name": "order_date", "type": "DATETIME", "description": "订单日期"}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
if os.path.exists(schema_path):
|
||||
with open(schema_path, 'r', encoding='utf-8') as f:
|
||||
self.table_schemas = json.load(f)
|
||||
else:
|
||||
self.table_schemas = default_schemas
|
||||
os.makedirs(data_dir, exist_ok=True)
|
||||
with open(schema_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(self.table_schemas, f, ensure_ascii=False, indent=2)
|
||||
|
||||
def vectorize_and_store(self):
|
||||
"""向量化数据并存储到Milvus"""
|
||||
self.create_collection()
|
||||
|
||||
all_texts = []
|
||||
all_metadata = []
|
||||
|
||||
for example in self.sql_examples:
|
||||
text = f"问题: {example['question']} SQL: {example['sql']} 描述: {example.get('description', '')}"
|
||||
all_texts.append(text)
|
||||
all_metadata.append({
|
||||
"content_type": "sql_example",
|
||||
"question": example['question'],
|
||||
"sql": example['sql'],
|
||||
"description": example.get('description', ''),
|
||||
"table_name": ""
|
||||
})
|
||||
|
||||
for schema in self.table_schemas:
|
||||
columns_desc = ", ".join([f"{col['name']} ({col['type']}): {col.get('description', '')}"
|
||||
for col in schema['columns']])
|
||||
text = f"表 {schema['table_name']}: {schema['description']} 字段: {columns_desc}"
|
||||
all_texts.append(text)
|
||||
all_metadata.append({
|
||||
"content_type": "table_schema",
|
||||
"question": "",
|
||||
"sql": "",
|
||||
"description": schema['description'],
|
||||
"table_name": schema['table_name']
|
||||
})
|
||||
|
||||
embeddings = self.embedding_function.encode_text(all_texts)
|
||||
|
||||
insert_data = []
|
||||
for i, (embedding, metadata) in enumerate(zip(embeddings, all_metadata)):
|
||||
insert_data.append([
|
||||
metadata["content_type"],
|
||||
metadata["question"],
|
||||
metadata["sql"],
|
||||
metadata["description"],
|
||||
metadata["table_name"],
|
||||
embedding.tolist()
|
||||
])
|
||||
|
||||
self.collection.insert(insert_data)
|
||||
self.collection.flush()
|
||||
self.collection.load()
|
||||
|
||||
def search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
|
||||
"""搜索相关的知识库信息"""
|
||||
if not self.data_loaded:
|
||||
self.load_data()
|
||||
|
||||
query_embedding = self.embedding_function.encode_text([query])[0]
|
||||
|
||||
search_params = {"metric_type": "IP", "params": {}}
|
||||
results = self.collection.search(
|
||||
[query_embedding.tolist()],
|
||||
anns_field="embedding",
|
||||
param=search_params,
|
||||
limit=top_k,
|
||||
output_fields=["content_type", "question", "sql", "description", "table_name"]
|
||||
)[0]
|
||||
|
||||
formatted_results = []
|
||||
for hit in results:
|
||||
result = {
|
||||
"score": float(hit.distance),
|
||||
"content_type": hit.entity.get("content_type"),
|
||||
"question": hit.entity.get("question"),
|
||||
"sql": hit.entity.get("sql"),
|
||||
"description": hit.entity.get("description"),
|
||||
"table_name": hit.entity.get("table_name")
|
||||
}
|
||||
formatted_results.append(result)
|
||||
|
||||
return formatted_results
|
||||
|
||||
def _fallback_search(self, query: str, top_k: int) -> List[Dict[str, Any]]:
|
||||
"""降级搜索方法（简单文本匹配）"""
|
||||
results = []
|
||||
query_lower = query.lower()
|
||||
|
||||
for example in self.sql_examples:
|
||||
question_lower = example['question'].lower()
|
||||
sql_lower = example['sql'].lower()
|
||||
|
||||
score = 0
|
||||
for word in query_lower.split():
|
||||
if word in question_lower:
|
||||
score += 2
|
||||
if word in sql_lower:
|
||||
score += 1
|
||||
|
||||
if score > 0:
|
||||
results.append({
|
||||
"score": score,
|
||||
"content_type": "sql_example",
|
||||
"question": example['question'],
|
||||
"sql": example['sql'],
|
||||
"description": example.get('description', ''),
|
||||
"table_name": ""
|
||||
})
|
||||
|
||||
results.sort(key=lambda x: x['score'], reverse=True)
|
||||
return results[:top_k]
|
||||
|
||||
def add_sql_example(self, question: str, sql: str, description: str = ""):
|
||||
"""添加新的SQL示例"""
|
||||
new_example = {
|
||||
"question": question,
|
||||
"sql": sql,
|
||||
"description": description,
|
||||
"database": "sqlite"
|
||||
}
|
||||
self.sql_examples.append(new_example)
|
||||
|
||||
data_dir = os.path.join(os.path.dirname(__file__), "data")
|
||||
sql_examples_path = os.path.join(data_dir, "qsql_examples.json")
|
||||
|
||||
with open(sql_examples_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(self.sql_examples, f, ensure_ascii=False, indent=2)
|
||||
|
||||
if self.collection and self.data_loaded:
|
||||
text = f"问题: {question} SQL: {sql} 描述: {description}"
|
||||
embedding = self.embedding_function.encode_text([text])[0]
|
||||
|
||||
insert_data = [[
|
||||
"sql_example",
|
||||
question,
|
||||
sql,
|
||||
description,
|
||||
"",
|
||||
embedding.tolist()
|
||||
]]
|
||||
|
||||
self.collection.insert(insert_data)
|
||||
self.collection.flush()
|
||||
|
||||
def cleanup(self):
|
||||
"""清理资源"""
|
||||
if self.collection:
|
||||
self.collection.release()
|
||||
|
||||
if self.milvus_client and self.milvus_client.has_collection(self.collection_name):
|
||||
self.milvus_client.drop_collection(self.collection_name)
|
||||
|
||||
|
||||
def demo():
|
||||
"""简单演示"""
|
||||
# 模型测试
|
||||
embedding_function = BGESmallEmbeddingFunction()
|
||||
test_texts = ["查询用户", "统计数据"]
|
||||
embeddings = embedding_function.encode_text(test_texts)
|
||||
print(f"向量维度: {embeddings.shape}")
|
||||
|
||||
# 数据库查询演示
|
||||
db_path = "demo.db"
|
||||
|
||||
if os.path.exists(db_path):
|
||||
os.remove(db_path)
|
||||
|
||||
conn = sqlite3.connect(db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute("CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT, age INTEGER, city TEXT)")
|
||||
|
||||
users_data = [(1, '张三', 25, '北京'), (2, '李四', 32, '上海'), (3, '王五', 35, '深圳')]
|
||||
cursor.executemany("INSERT INTO users VALUES (?, ?, ?, ?)", users_data)
|
||||
|
||||
conn.commit()
|
||||
|
||||
# 执行查询
|
||||
test_sqls = [
|
||||
("查询所有用户", "SELECT * FROM users"),
|
||||
("年龄大于30的用户", "SELECT * FROM users WHERE age > 30"),
|
||||
("统计用户总数", "SELECT COUNT(*) FROM users")
|
||||
]
|
||||
|
||||
for i, (question, sql) in enumerate(test_sqls, 1):
|
||||
print(f"\n问题 {i}: {question}")
|
||||
print("-" * 40)
|
||||
print(f"SQL: {sql}")
|
||||
|
||||
cursor.execute(sql)
|
||||
rows = cursor.fetchall()
|
||||
|
||||
if rows:
|
||||
print(f"返回 {len(rows)} 行数据")
|
||||
for j, row in enumerate(rows[:2], 1):
|
||||
print(f" {j}. {row}")
|
||||
|
||||
if len(rows) > 2:
|
||||
print(f" ... 还有 {len(rows) - 2} 行")
|
||||
else:
|
||||
print("无数据返回")
|
||||
|
||||
conn.close()
|
||||
os.remove(db_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
demo()
|
||||
Reference in New Issue
Block a user