mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-09 01:36:57 +08:00
92 lines
3.2 KiB
Python
92 lines
3.2 KiB
Python
import json
|
|
import time
|
|
import traceback
|
|
from typing import List
|
|
|
|
from sqlalchemy import and_, select, update
|
|
from sqlalchemy.orm.session import Session
|
|
|
|
from apps.ai_model.embedding import EmbeddingModelCache
|
|
from common.core.config import settings
|
|
from common.core.deps import SessionDep
|
|
from common.utils.utils import SQLBotLogUtil
|
|
from ..models.datasource import CoreTable, CoreField
|
|
|
|
|
|
def delete_table_by_ds_id(session: SessionDep, id: int):
|
|
session.query(CoreTable).filter(CoreTable.ds_id == id).delete(synchronize_session=False)
|
|
session.commit()
|
|
|
|
|
|
def get_tables_by_ds_id(session: SessionDep, id: int):
|
|
return session.query(CoreTable).filter(CoreTable.ds_id == id).order_by(
|
|
CoreTable.table_name.asc()).all()
|
|
|
|
|
|
def update_table(session: SessionDep, item: CoreTable):
|
|
record = session.query(CoreTable).filter(CoreTable.id == item.id).first()
|
|
record.checked = item.checked
|
|
record.custom_comment = item.custom_comment
|
|
session.add(record)
|
|
session.commit()
|
|
|
|
|
|
def run_fill_empty_table_embedding(session: Session):
|
|
if not settings.EMBEDDING_ENABLED:
|
|
return
|
|
|
|
SQLBotLogUtil.info('get tables')
|
|
stmt = select(CoreTable.id).where(and_(CoreTable.embedding.is_(None)))
|
|
results = session.execute(stmt).scalars().all()
|
|
|
|
save_table_embedding(session, results)
|
|
|
|
|
|
def save_table_embedding(session: Session, ids: List[int]):
|
|
if not settings.EMBEDDING_ENABLED:
|
|
return
|
|
|
|
if not ids or len(ids) == 0:
|
|
return
|
|
try:
|
|
SQLBotLogUtil.info('start table embedding')
|
|
start_time = time.time()
|
|
model = EmbeddingModelCache.get_model()
|
|
for _id in ids:
|
|
table = session.query(CoreTable).filter(CoreTable.id == _id).first()
|
|
fields = session.query(CoreField).filter(CoreField.table_id == table.id).all()
|
|
|
|
schema_table = ''
|
|
schema_table += f"# Table: {table.table_name}"
|
|
table_comment = ''
|
|
if table.custom_comment:
|
|
table_comment = table.custom_comment.strip()
|
|
if table_comment == '':
|
|
schema_table += '\n[\n'
|
|
else:
|
|
schema_table += f", {table_comment}\n[\n"
|
|
|
|
if fields:
|
|
field_list = []
|
|
for field in fields:
|
|
field_comment = ''
|
|
if field.custom_comment:
|
|
field_comment = field.custom_comment.strip()
|
|
if field_comment == '':
|
|
field_list.append(f"({field.field_name}:{field.field_type})")
|
|
else:
|
|
field_list.append(f"({field.field_name}:{field.field_type}, {field_comment})")
|
|
schema_table += ",\n".join(field_list)
|
|
schema_table += '\n]\n'
|
|
# table_schema.append(schema_table)
|
|
emb = json.dumps(model.embed_query(schema_table))
|
|
|
|
stmt = update(CoreTable).where(and_(CoreTable.id == _id)).values(embedding=emb)
|
|
session.execute(stmt)
|
|
session.commit()
|
|
|
|
end_time = time.time()
|
|
SQLBotLogUtil.info('table embedding finished in:' + str(end_time - start_time) + 'seconds')
|
|
except Exception:
|
|
traceback.print_exc()
|