feat: add SQL examples in prompt

This commit is contained in:
ulleo
2025-09-16 18:08:53 +08:00
parent 6f783512e5
commit f86b2f60cf
19 changed files with 506 additions and 31 deletions
+1
View File
@@ -26,6 +26,7 @@ config = context.config
# from apps.settings.models.setting_models import SQLModel
# from apps.chat.models.chat_model import SQLModel
from apps.terminology.models.terminology_model import SQLModel
from apps.data_training.models.data_training_model import SQLModel
# from apps.dashboard.models.dashboard_model import SQLModel
from common.core.config import settings # noqa
@@ -0,0 +1,41 @@
"""042_data_training
Revision ID: a487d9c69341
Revises: c4c3c36b720d
Create Date: 2025-09-15 15:41:43.332771
"""
from alembic import op
import sqlalchemy as sa
import sqlmodel.sql.sqltypes
from sqlalchemy.dialects import postgresql
import pgvector
# revision identifiers, used by Alembic.
revision = 'a487d9c69341'
down_revision = 'c4c3c36b720d'
branch_labels = None
depends_on = None
def upgrade():
# ### commands auto generated by Alembic - please adjust! ###
op.create_table('data_training',
sa.Column('id', sa.BigInteger(), sa.Identity(always=True), nullable=False),
sa.Column('oid', sa.BigInteger(), nullable=True),
sa.Column('datasource', sa.BigInteger(), nullable=True),
sa.Column('create_time', sa.DateTime(), nullable=True),
sa.Column('question', sqlmodel.sql.sqltypes.AutoString(length=255), nullable=True),
sa.Column('description', sa.Text(), nullable=True),
sa.Column('embedding', pgvector.sqlalchemy.vector.VECTOR(), nullable=True),
sa.PrimaryKeyConstraint('id')
)
# ### end Alembic commands ###
def downgrade():
# ### commands auto generated by Alembic - please adjust! ###
op.drop_table('data_training')
# ### end Alembic commands ###
+2
View File
@@ -1,6 +1,7 @@
from fastapi import APIRouter
from apps.terminology.api import terminology
from apps.data_training.api import data_training
from apps.chat.api import chat
from apps.dashboard.api import dashboard_api
from apps.datasource.api import datasource
@@ -14,6 +15,7 @@ api_router.include_router(workspace.router)
api_router.include_router(assistant.router)
api_router.include_router(aimodel.router)
api_router.include_router(terminology.router)
api_router.include_router(data_training.router)
api_router.include_router(datasource.router)
api_router.include_router(chat.router)
api_router.include_router(dashboard_api.router)
+3 -1
View File
@@ -172,11 +172,13 @@ class AiModelQuestion(BaseModel):
filter: str = []
sub_query: Optional[list[dict]] = None
terminologies: str = ""
data_training: str = ""
error_msg: str = ""
def sql_sys_question(self):
return get_sql_template()['system'].format(engine=self.engine, schema=self.db_schema, question=self.question,
lang=self.lang, terminologies=self.terminologies)
lang=self.lang, terminologies=self.terminologies,
data_training=self.data_training)
def sql_user_question(self, current_time: str):
return get_sql_template()['user'].format(engine=self.engine, schema=self.db_schema, question=self.question,
@@ -0,0 +1,39 @@
from typing import Optional
from fastapi import APIRouter, Query
from apps.data_training.curd.data_training import page_data_training, create_training, update_training, delete_training
from apps.data_training.models.data_training_model import DataTrainingInfo
from common.core.deps import SessionDep, CurrentUser, Trans
router = APIRouter(tags=["DataTraining"], prefix="/system/data-training")
@router.get("/page/{current_page}/{page_size}")
async def pager(session: SessionDep, current_user: CurrentUser, current_page: int, page_size: int,
question: Optional[str] = Query(None, description="搜索问题(可选)")):
current_page, page_size, total_count, total_pages, _list = page_data_training(session, current_page, page_size,
question,
current_user.oid)
return {
"current_page": current_page,
"page_size": page_size,
"total_count": total_count,
"total_pages": total_pages,
"data": _list
}
@router.put("")
async def create_or_update(session: SessionDep, current_user: CurrentUser, trans: Trans, info: DataTrainingInfo):
oid = current_user.oid
if info.id:
return update_training(session, info, oid, trans)
else:
return create_training(session, info, oid, trans)
@router.delete("")
async def delete(session: SessionDep, id_list: list[int]):
delete_training(session, id_list)
@@ -0,0 +1,312 @@
import datetime
import logging
import traceback
from typing import List, Optional
from xml.dom.minidom import parseString
import dicttoxml
from sqlalchemy import and_, select, func, delete, update, or_
from sqlalchemy import text
from sqlalchemy.orm.session import Session
from apps.ai_model.embedding import EmbeddingModelCache
from apps.data_training.models.data_training_model import DataTrainingInfo, DataTraining
from apps.datasource.models.datasource import CoreDatasource
from apps.template.generate_chart.generator import get_base_data_training_template
from common.core.config import settings
from common.core.deps import SessionDep, Trans
from common.utils.embedding_threads import run_save_data_training_embeddings
def page_data_training(session: SessionDep, current_page: int = 1, page_size: int = 10, name: Optional[str] = None,
oid: Optional[int] = 1):
_list: List[DataTrainingInfo] = []
current_page = max(1, current_page)
page_size = max(10, page_size)
total_count = 0
total_pages = 0
if name and name.strip() != "":
keyword_pattern = f"%{name.strip()}%"
parent_ids_subquery = (
select(DataTraining.id)
.where(and_(DataTraining.question.ilike(keyword_pattern), DataTraining.oid == oid)) # LIKE查询条件
)
else:
parent_ids_subquery = (
select(DataTraining.id).where(and_(DataTraining.oid == oid))
)
count_stmt = select(func.count()).select_from(parent_ids_subquery.subquery())
total_count = session.execute(count_stmt).scalar()
total_pages = (total_count + page_size - 1) // page_size
if current_page > total_pages:
current_page = 1
paginated_parent_ids = (
parent_ids_subquery
.order_by(DataTraining.create_time.desc())
.offset((current_page - 1) * page_size)
.limit(page_size)
.subquery()
)
stmt = (
select(
DataTraining.id,
DataTraining.oid,
DataTraining.datasource,
CoreDatasource.name,
DataTraining.question,
DataTraining.create_time,
DataTraining.description,
)
.outerjoin(CoreDatasource, and_(DataTraining.datasource == CoreDatasource.id))
.where(and_(DataTraining.id.in_(paginated_parent_ids)))
.order_by(DataTraining.create_time.desc())
)
result = session.execute(stmt)
for row in result:
_list.append(DataTrainingInfo(
id=row.id,
oid=row.oid,
datasource=row.datasource,
datasource_name=row.name,
question=row.question,
create_time=row.create_time,
description=row.description,
))
return current_page, page_size, total_count, total_pages, _list
def create_training(session: SessionDep, info: DataTrainingInfo, oid: int, trans: Trans):
create_time = datetime.datetime.now()
if info.datasource is None:
raise Exception(trans("i18n_data_training.datasource_cannot_be_none"))
parent = DataTraining(question=info.question, create_time=create_time, description=info.description, oid=oid,
datasource=info.datasource)
exists = session.query(
session.query(DataTraining).filter(
and_(DataTraining.question == info.question, DataTraining.oid == oid)).exists()).scalar()
if exists:
raise Exception(trans("i18n_data_training.exists_in_db"))
result = DataTraining(**parent.model_dump())
session.add(parent)
session.flush()
session.refresh(parent)
result.id = parent.id
session.commit()
# embedding
run_save_data_training_embeddings([result.id])
return result.id
def update_training(session: SessionDep, info: DataTrainingInfo, oid: int, trans: Trans):
count = session.query(DataTraining).filter(
DataTraining.oid == oid,
DataTraining.id == info.id
).count()
if count == 0:
raise Exception(trans('i18n_data_training.data_training_not_exists'))
exists = session.query(
session.query(DataTraining).filter(
and_(DataTraining.question == info.question, DataTraining.oid == oid,
DataTraining.id != info.id)).exists()).scalar()
if exists:
raise Exception(trans("i18n_data_training.exists_in_db"))
stmt = update(DataTraining).where(and_(DataTraining.id == info.id)).values(
question=info.question,
description=info.description,
)
session.execute(stmt)
session.commit()
# embedding
run_save_data_training_embeddings([info.id])
return info.id
def delete_training(session: SessionDep, ids: list[int]):
stmt = delete(DataTraining).where(and_(DataTraining.id.in_(ids)))
session.execute(stmt)
session.commit()
# def run_save_embeddings(ids: List[int]):
# executor.submit(save_embeddings, ids)
#
#
# def fill_empty_embeddings():
# executor.submit(run_fill_empty_embeddings)
def run_fill_empty_embeddings(session: Session):
if not settings.EMBEDDING_ENABLED:
return
stmt = select(DataTraining.id).where(and_(DataTraining.embedding.is_(None)))
results = session.execute(stmt).scalars().all()
save_embeddings(session, results)
def save_embeddings(session: Session, ids: List[int]):
if not settings.EMBEDDING_ENABLED:
return
if not ids or len(ids) == 0:
return
try:
_list = session.query(DataTraining).filter(and_(DataTraining.id.in_(ids))).all()
_question_list = [item.question for item in _list]
model = EmbeddingModelCache.get_model()
results = model.embed_documents(_question_list)
for index in range(len(results)):
item = results[index]
stmt = update(DataTraining).where(and_(DataTraining.id == _list[index].id)).values(embedding=item)
session.execute(stmt)
session.commit()
except Exception:
traceback.print_exc()
embedding_sql = f"""
SELECT id, datasource, question, similarity
FROM
(SELECT id, datasource, question, oid,
( 1 - (embedding <=> :embedding_array) ) AS similarity
FROM data_training AS child
) TEMP
WHERE similarity > {settings.EMBEDDING_SIMILARITY} and oid = :oid and datasource = :datasource
ORDER BY similarity DESC
LIMIT {settings.EMBEDDING_DATA_TRAINING_TOP_COUNT}
"""
def select_terminology_by_question(session: SessionDep, question: str, oid: int, datasource: int):
if question.strip() == "":
return []
_list: List[DataTraining] = []
# maybe use label later?
stmt = (
select(
DataTraining.id,
DataTraining.question,
)
.where(
and_(or_(text(":sentence ILIKE '%' || question || '%'"), text("question ILIKE '%' || :sentence || '%'")),
DataTraining.oid == oid,
DataTraining.datasource == datasource)
)
)
results = session.execute(stmt, {'sentence': question}).fetchall()
for row in results:
_list.append(DataTraining(id=row.id, question=row.question))
if settings.EMBEDDING_ENABLED:
try:
model = EmbeddingModelCache.get_model()
embedding = model.embed_query(question)
results = session.execute(text(embedding_sql),
{'embedding_array': str(embedding), 'oid': oid, 'datasource': datasource})
for row in results:
_list.append(DataTraining(id=row.id, question=row.question))
except Exception:
traceback.print_exc()
_map: dict = {}
_ids: list[int] = []
for row in _list:
if row.id in _ids:
continue
else:
_ids.append(row.id)
if len(_ids) == 0:
return []
t_list = session.query(DataTraining.id, DataTraining.datasource, DataTraining.question,
DataTraining.description).filter(
and_(DataTraining.id.in_(_ids))).all()
for row in t_list:
_map[row.id] = {'question': row.question, 'suggestion-answer': row.description}
_results: list[dict] = []
for key in _map.keys():
_results.append(_map.get(key))
return _results
def to_xml_string(_dict: list[dict] | dict, root: str = 'sql-examples') -> str:
item_name_func = lambda x: 'sql-example' if x == 'sql-examples' else 'item'
dicttoxml.LOG.setLevel(logging.ERROR)
xml = dicttoxml.dicttoxml(_dict,
cdata=['question', 'suggestion-answer'],
custom_root=root,
item_func=item_name_func,
xml_declaration=False,
encoding='utf-8',
attr_type=False).decode('utf-8')
pretty_xml = parseString(xml).toprettyxml()
if pretty_xml.startswith('<?xml'):
end_index = pretty_xml.find('>') + 1
pretty_xml = pretty_xml[end_index:].lstrip()
# 替换所有 XML 转义字符
escape_map = {
'&lt;': '<',
'&gt;': '>',
'&amp;': '&',
'&quot;': '"',
'&apos;': "'"
}
for escaped, original in escape_map.items():
pretty_xml = pretty_xml.replace(escaped, original)
return pretty_xml
def get_training_template(session: SessionDep, question: str, datasource: int, oid: Optional[int] = 1) -> str:
if not oid:
oid = 1
if not datasource:
return ''
_results = select_terminology_by_question(session, question, oid, datasource)
if _results and len(_results) > 0:
data_training = to_xml_string(_results)
template = get_base_data_training_template().format(data_training=data_training)
return template
else:
return ''
@@ -0,0 +1,28 @@
from datetime import datetime
from typing import List, Optional
from pgvector.sqlalchemy import VECTOR
from pydantic import BaseModel
from sqlalchemy import Column, Text, BigInteger, DateTime, Identity
from sqlmodel import SQLModel, Field
class DataTraining(SQLModel, table=True):
__tablename__ = "data_training"
id: Optional[int] = Field(sa_column=Column(BigInteger, Identity(always=True), primary_key=True))
oid: Optional[int] = Field(sa_column=Column(BigInteger, nullable=True, default=1))
datasource: Optional[int] = Field(sa_column=Column(BigInteger, nullable=True))
create_time: Optional[datetime] = Field(sa_column=Column(DateTime(timezone=False), nullable=True))
question: Optional[str] = Field(max_length=255)
description: Optional[str] = Field(sa_column=Column(Text, nullable=True))
embedding: Optional[List[float]] = Field(sa_column=Column(VECTOR(), nullable=True))
class DataTrainingInfo(BaseModel):
id: Optional[int] = None
oid: Optional[int] = None
datasource: Optional[int] = None
datasource_name: Optional[str] = None
create_time: Optional[datetime] = None
question: Optional[str] = None
description: Optional[str] = None
@@ -8,3 +8,7 @@ def get_chart_template():
def get_base_terminology_template():
template = get_base_template()
return template['template']['terminology']
def get_base_data_training_template():
template = get_base_template()
return template['template']['data_training']
+15 -23
View File
@@ -1,23 +1,21 @@
import datetime
import logging
import traceback
from concurrent.futures import ThreadPoolExecutor
from typing import List, Optional
from xml.dom.minidom import parseString
import dicttoxml
from sqlalchemy import and_, or_, select, func, delete, update, union
from sqlalchemy import create_engine, text
from sqlalchemy import text
from sqlalchemy.orm import aliased
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm.session import Session
from apps.ai_model.embedding import EmbeddingModelCache
from apps.template.generate_chart.generator import get_base_terminology_template
from apps.terminology.models.terminology_model import Terminology, TerminologyInfo
from common.core.config import settings
from common.core.deps import SessionDep, Trans
executor = ThreadPoolExecutor(max_workers=200)
from common.utils.embedding_threads import run_save_terminology_embeddings
def page_terminology(session: SessionDep, current_page: int = 1, page_size: int = 10, name: Optional[str] = None,
@@ -183,7 +181,7 @@ def create_terminology(session: SessionDep, info: TerminologyInfo, oid: int, tra
session.commit()
# embedding
run_save_embeddings([result.id])
run_save_terminology_embeddings([result.id])
return result.id
@@ -240,7 +238,7 @@ def update_terminology(session: SessionDep, info: TerminologyInfo, oid: int, tra
session.commit()
# embedding
run_save_embeddings([info.id])
run_save_terminology_embeddings([info.id])
return info.id
@@ -251,37 +249,31 @@ def delete_terminology(session: SessionDep, ids: list[int]):
session.commit()
def run_save_embeddings(ids: List[int]):
executor.submit(save_embeddings, ids)
# def run_save_embeddings(ids: List[int]):
# executor.submit(save_embeddings, ids)
#
#
# def fill_empty_embeddings():
# executor.submit(run_fill_empty_embeddings)
def fill_empty_embeddings():
executor.submit(run_fill_empty_embeddings)
def run_fill_empty_embeddings():
def run_fill_empty_embeddings(session: Session):
if not settings.EMBEDDING_ENABLED:
return
engine = create_engine(str(settings.SQLALCHEMY_DATABASE_URI))
session_maker = sessionmaker(bind=engine)
session = session_maker()
stmt1 = select(Terminology.id).where(and_(Terminology.embedding.is_(None), Terminology.pid.is_(None)))
stmt2 = select(Terminology.pid).where(and_(Terminology.embedding.is_(None), Terminology.pid.isnot(None))).distinct()
combined_stmt = union(stmt1, stmt2)
results = session.execute(combined_stmt).scalars().all()
save_embeddings(results)
save_embeddings(session, results)
def save_embeddings(ids: List[int]):
def save_embeddings(session: Session, ids: List[int]):
if not settings.EMBEDDING_ENABLED:
return
if not ids or len(ids) == 0:
return
try:
engine = create_engine(str(settings.SQLALCHEMY_DATABASE_URI))
session_maker = sessionmaker(bind=engine)
session = session_maker()
_list = session.query(Terminology).filter(or_(Terminology.id.in_(ids), Terminology.pid.in_(ids))).all()
@@ -310,7 +302,7 @@ FROM terminology AS child
) TEMP
WHERE similarity > {settings.EMBEDDING_SIMILARITY} and oid = :oid
ORDER BY similarity DESC
LIMIT {settings.EMBEDDING_TOP_COUNT}
LIMIT {settings.EMBEDDING_TERMINOLOGY_TOP_COUNT}
"""
+3 -1
View File
@@ -90,7 +90,9 @@ class Settings(BaseSettings):
DEFAULT_EMBEDDING_MODEL: str = 'shibing624/text2vec-base-chinese'
EMBEDDING_ENABLED: bool = True
EMBEDDING_SIMILARITY: float = 0.4
EMBEDDING_TOP_COUNT: int = 5
EMBEDDING_DEFAULT_TOP_COUNT: int = 5
EMBEDDING_TERMINOLOGY_TOP_COUNT: int = EMBEDDING_DEFAULT_TOP_COUNT
EMBEDDING_DATA_TRAINING_TOP_COUNT: int = EMBEDDING_DEFAULT_TOP_COUNT
PG_POOL_SIZE: int = 20
PG_MAX_OVERFLOW: int = 30
+33
View File
@@ -0,0 +1,33 @@
from concurrent.futures import ThreadPoolExecutor
from typing import List
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from common.core.config import settings
executor = ThreadPoolExecutor(max_workers=200)
engine = create_engine(str(settings.SQLALCHEMY_DATABASE_URI))
session_maker = sessionmaker(bind=engine)
session = session_maker()
def run_save_terminology_embeddings(ids: List[int]):
from apps.terminology.curd.terminology import save_embeddings
executor.submit(save_embeddings, session, ids)
def fill_empty_terminology_embeddings():
from apps.terminology.curd.terminology import run_fill_empty_embeddings
executor.submit(run_fill_empty_embeddings, session)
def run_save_data_training_embeddings(ids: List[int]):
from apps.data_training.curd.data_training import save_embeddings
executor.submit(save_embeddings, session, ids)
def fill_empty_data_training_embeddings():
from apps.data_training.curd.data_training import run_fill_empty_embeddings
executor.submit(run_fill_empty_embeddings, session)
+5
View File
@@ -42,5 +42,10 @@
"terminology_not_exists": "Terminology does not exists",
"cannot_be_repeated": "Term name, synonyms cannot be repeated",
"exists_in_db": "Term name, synonyms exists"
},
"i18n_data_training": {
"datasource_cannot_be_none": "Datasource cannot be none",
"data_training_not_exists": "Example does not exists",
"exists_in_db": "Question exists"
}
}
+6 -1
View File
@@ -41,6 +41,11 @@
"i18n_terminology": {
"terminology_not_exists": "该术语不存在",
"cannot_be_repeated": "术语名称,同义词不能重复",
"exists_in_db": "术语名称,同义词已存在"
"exists_in_db": "术语名称,同义词已存在"
},
"i18n_data_training": {
"datasource_cannot_be_none": "数据源不能为空",
"data_training_not_exists": "该示例不存在",
"exists_in_db": "该问题已存在"
}
}
+9 -4
View File
@@ -15,10 +15,10 @@ from apps.api import api_router
from apps.system.crud.aimodel_manage import async_model_info
from apps.system.crud.assistant import init_dynamic_cors
from apps.system.middleware.auth import TokenMiddleware
from apps.terminology.curd.terminology import fill_empty_embeddings
from common.core.config import settings
from common.core.response_middleware import ResponseMiddleware, exception_handler
from common.core.sqlbot_cache import init_sqlbot_cache
from common.utils.embedding_threads import fill_empty_terminology_embeddings, fill_empty_data_training_embeddings
from common.utils.utils import SQLBotLogUtil
@@ -27,8 +27,12 @@ def run_migrations():
command.upgrade(alembic_cfg, "head")
def init_embedding_data():
fill_empty_embeddings()
def init_terminology_embedding_data():
fill_empty_terminology_embeddings()
def init_data_training_embedding_data():
fill_empty_data_training_embeddings()
@asynccontextmanager
@@ -36,7 +40,8 @@ async def lifespan(app: FastAPI):
run_migrations()
init_sqlbot_cache()
init_dynamic_cors(app)
init_embedding_data()
init_terminology_embedding_data()
init_data_training_embedding_data()
SQLBotLogUtil.info("✅ SQLBot 初始化完成")
await sqlbot_xpack.core.clean_xpack_cache()
await async_model_info() # 异步加密已有模型的密钥和地址
+5 -1
View File
@@ -2,7 +2,9 @@ template:
terminology: |
{terminologies}
data_training: |
{data_training}
sql:
system: |
<Instruction>
@@ -13,6 +15,7 @@ template:
其中,<db-engine>:提供数据库引擎及版本信息;
<m-schema>:以 M-Schema 格式提供数据库表结构信息;
<terminologies>:提供一组术语,块内每一个<terminology>就是术语,其中同一个<words>内的多个<word>代表术语的多种叫法,也就是术语与它的同义词,<description>即该术语对应的描述,其中也可能是能够用来参考的计算公式,或者是一些其他的查询条件
<sql-examples>:提供一组SQL示例,你可以参考这些示例来生成你的回答,其中<question>内是提问,<suggestion-answer>内是对于该<question>提问的解释或者对应应该回答的SQL示例
用户的提问在<user-question>内,<error-msg>内则会提供上次执行你提供的SQL时会出现的错误信息,<background-infos>内的<current-time>会告诉你用户当前提问的时间
</Instruction>
@@ -224,6 +227,7 @@ template:
</m-schema>
{terminologies}
{data_training}
</Info>
### 响应, 请根据上述要求直接返回JSON结果: