| 55 | |
| 56 | |
| 57 | def save_table_embedding(session_maker, ids: List[int]): |
| 58 | if not settings.TABLE_EMBEDDING_ENABLED: |
| 59 | return |
| 60 | |
| 61 | if not ids or len(ids) == 0: |
| 62 | return |
| 63 | try: |
| 64 | SQLBotLogUtil.info('start table embedding') |
| 65 | start_time = time.time() |
| 66 | model = EmbeddingModelCache.get_model() |
| 67 | session = session_maker() |
| 68 | for _id in ids: |
| 69 | table = session.query(CoreTable).filter(CoreTable.id == _id).first() |
| 70 | fields = session.query(CoreField).filter(CoreField.table_id == table.id).all() |
| 71 | |
| 72 | schema_table = '' |
| 73 | schema_table += f"# Table: {table.table_name}" |
| 74 | table_comment = '' |
| 75 | if table.custom_comment: |
| 76 | table_comment = table.custom_comment.strip() |
| 77 | if table_comment == '': |
| 78 | schema_table += '\n[\n' |
| 79 | else: |
| 80 | schema_table += f", {table_comment}\n[\n" |
| 81 | |
| 82 | if fields: |
| 83 | field_list = [] |
| 84 | for field in fields: |
| 85 | field_comment = '' |
| 86 | if field.custom_comment: |
| 87 | field_comment = field.custom_comment.strip() |
| 88 | if field_comment == '': |
| 89 | field_list.append(f"({field.field_name}:{field.field_type})") |
| 90 | else: |
| 91 | field_list.append(f"({field.field_name}:{field.field_type}, {field_comment})") |
| 92 | schema_table += ",\n".join(field_list) |
| 93 | schema_table += '\n]\n' |
| 94 | # table_schema.append(schema_table) |
| 95 | emb = json.dumps(model.embed_query(schema_table)) |
| 96 | |
| 97 | stmt = update(CoreTable).where(and_(CoreTable.id == _id)).values(embedding=emb) |
| 98 | session.execute(stmt) |
| 99 | session.commit() |
| 100 | |
| 101 | end_time = time.time() |
| 102 | SQLBotLogUtil.info('table embedding finished in: ' + str(end_time - start_time) + ' seconds') |
| 103 | except Exception: |
| 104 | traceback.print_exc() |
| 105 | finally: |
| 106 | session_maker.remove() |
| 107 | |
| 108 | |
| 109 | def save_ds_embedding(session_maker, ids: List[int]): |