mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
feat: add SQL examples in prompt
This commit is contained in:
@@ -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 ###
|
||||
@@ -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)
|
||||
|
||||
@@ -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 = {
|
||||
'<': '<',
|
||||
'>': '>',
|
||||
'&': '&',
|
||||
'"': '"',
|
||||
''': "'"
|
||||
}
|
||||
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']
|
||||
|
||||
@@ -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}
|
||||
"""
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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() # 异步加密已有模型的密钥和地址
|
||||
|
||||
@@ -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结果:
|
||||
|
||||
Reference in New Issue
Block a user