| 107 | |
| 108 | |
| 109 | def save_ds_embedding(session_maker, ids: List[int]): |
| 110 | if not settings.TABLE_EMBEDDING_ENABLED: |
| 111 | return |
| 112 | |
| 113 | if not ids or len(ids) == 0: |
| 114 | return |
| 115 | try: |
| 116 | SQLBotLogUtil.info('start datasource embedding') |
| 117 | start_time = time.time() |
| 118 | model = EmbeddingModelCache.get_model() |
| 119 | session = session_maker() |
| 120 | for _id in ids: |
| 121 | schema_table = '' |
| 122 | ds = session.query(CoreDatasource).filter(CoreDatasource.id == _id).first() |
| 123 | schema_table += f"{ds.name}, {ds.description}\n" |
| 124 | tables = session.query(CoreTable).filter(CoreTable.ds_id == ds.id).all() |
| 125 | for table in tables: |
| 126 | fields = session.query(CoreField).filter(CoreField.table_id == table.id).all() |
| 127 | |
| 128 | schema_table += f"# Table: {table.table_name}" |
| 129 | table_comment = '' |
| 130 | if table.custom_comment: |
| 131 | table_comment = table.custom_comment.strip() |
| 132 | if table_comment == '': |
| 133 | schema_table += '\n[\n' |
| 134 | else: |
| 135 | schema_table += f", {table_comment}\n[\n" |
| 136 | |
| 137 | if fields: |
| 138 | field_list = [] |
| 139 | for field in fields: |
| 140 | field_comment = '' |
| 141 | if field.custom_comment: |
| 142 | field_comment = field.custom_comment.strip() |
| 143 | if field_comment == '': |
| 144 | field_list.append(f"({field.field_name}:{field.field_type})") |
| 145 | else: |
| 146 | field_list.append(f"({field.field_name}:{field.field_type}, {field_comment})") |
| 147 | schema_table += ",\n".join(field_list) |
| 148 | schema_table += '\n]\n' |
| 149 | # table_schema.append(schema_table) |
| 150 | emb = json.dumps(model.embed_query(schema_table)) |
| 151 | |
| 152 | stmt = update(CoreDatasource).where(and_(CoreDatasource.id == _id)).values(embedding=emb) |
| 153 | session.execute(stmt) |
| 154 | session.commit() |
| 155 | |
| 156 | end_time = time.time() |
| 157 | SQLBotLogUtil.info('datasource embedding finished in: ' + str(end_time - start_time) + ' seconds') |
| 158 | except Exception: |
| 159 | traceback.print_exc() |
| 160 | finally: |
| 161 | session_maker.remove() |