(session: SessionDep, current_user: CurrentUser, ds: CoreDatasource, question: str,
embedding: bool = True, table_list: list[str] = None)
| 506 | |
| 507 | |
| 508 | def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDatasource, question: str, |
| 509 | embedding: bool = True, table_list: list[str] = None) -> tuple[str, list]: |
| 510 | schema_str = "" |
| 511 | table_objs = get_table_obj_by_ds(session=session, current_user=current_user, ds=ds) |
| 512 | if len(table_objs) == 0: |
| 513 | return schema_str, [] |
| 514 | db_name = table_objs[0].schema |
| 515 | schema_str += f"【DB_ID】 {db_name}\n【Schema】\n" |
| 516 | tables = [] |
| 517 | all_tables = [] # temp save all tables |
| 518 | table_name_list = [] |
| 519 | for obj in table_objs: |
| 520 | # 如果传入了table_list，则只处理在列表中的表 |
| 521 | if table_list is not None and obj.table.table_name not in table_list: |
| 522 | continue |
| 523 | |
| 524 | schema_table = '' |
| 525 | no_schema_types = ["mysql", "es", "sqlite", "hive", "doris", "starrocks"] |
| 526 | schema_table += f"# Table: {db_name}.{obj.table.table_name}" if ds.type not in no_schema_types and db_name else f"# Table: {obj.table.table_name}" |
| 527 | table_comment = '' |
| 528 | if obj.table.custom_comment: |
| 529 | table_comment = obj.table.custom_comment.strip() |
| 530 | if table_comment == '': |
| 531 | schema_table += '\n[\n' |
| 532 | else: |
| 533 | schema_table += f", {table_comment}\n[\n" |
| 534 | |
| 535 | if obj.fields: |
| 536 | field_list = [] |
| 537 | for field in obj.fields: |
| 538 | field_comment = '' |
| 539 | if field.custom_comment: |
| 540 | field_comment = field.custom_comment.strip() |
| 541 | if field_comment == '': |
| 542 | field_list.append(f"({field.field_name}:{field.field_type})") |
| 543 | else: |
| 544 | field_list.append(f"({field.field_name}:{field.field_type}, {field_comment})") |
| 545 | schema_table += ",\n".join(field_list) |
| 546 | schema_table += '\n]\n' |
| 547 | |
| 548 | t_obj = {"id": obj.table.id, "table_name": obj.table.table_name, "schema_table": schema_table, |
| 549 | "embedding": obj.table.embedding} |
| 550 | tables.append(t_obj) |
| 551 | all_tables.append(t_obj) |
| 552 | |
| 553 | # 如果没有符合过滤条件的表，直接返回 |
| 554 | if not tables: |
| 555 | return schema_str, [] |
| 556 | |
| 557 | # do table embedding |
| 558 | if embedding and tables and settings.TABLE_EMBEDDING_ENABLED: |
| 559 | tables = calc_table_embedding(tables, question) |
| 560 | # splice schema |
| 561 | if tables: |
| 562 | for s in tables: |
| 563 | schema_str += s.get('schema_table') |
| 564 | table_name_list.append(s.get('table_name')) |
| 565 |
no test coverage detected