From 090df35b4880d0858dd7a119cce512e69be0dfa6 Mon Sep 17 00:00:00 2001 From: junjun Date: Thu, 9 Oct 2025 12:51:32 +0800 Subject: [PATCH] feat: Vector retrieval matches tables --- backend/apps/datasource/crud/table.py | 36 +++++++++++---------------- 1 file changed, 14 insertions(+), 22 deletions(-) diff --git a/backend/apps/datasource/crud/table.py b/backend/apps/datasource/crud/table.py index b7bb4ab0..b736d34a 100644 --- a/backend/apps/datasource/crud/table.py +++ b/backend/apps/datasource/crud/table.py @@ -48,18 +48,18 @@ def save_table_embedding(session: SessionDep, ids: List[int]): if not ids or len(ids) == 0: return try: - - _list = session.query(CoreTable).filter(and_(CoreTable.id.in_(ids))).all() - - table_schema = [] - for item in _list: - fields = session.query(CoreField).filter(CoreField.table_id == item.id).all() + 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: {item.table_name}" + schema_table += f"# Table: {table.table_name}" table_comment = '' - if item.custom_comment: - table_comment = item.custom_comment.strip() + if table.custom_comment: + table_comment = table.custom_comment.strip() if table_comment == '': schema_table += '\n[\n' else: @@ -77,22 +77,14 @@ def save_table_embedding(session: SessionDep, ids: List[int]): 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) + # table_schema.append(schema_table) + emb = model.embed_query(schema_table) - model = EmbeddingModelCache.get_model() - - SQLBotLogUtil.info(json.dumps(table_schema)) - SQLBotLogUtil.info('start table embedding') - start_time = time.time() - results = model.embed_documents(table_schema) - end_time = time.time() - SQLBotLogUtil.info('table embedding finished in:' + str(end_time - start_time) + 'seconds') - - for index in range(len(results)): - item = results[index] - stmt = update(CoreTable).where(and_(CoreTable.id == _list[index].id)).values(embedding=item) + 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()