浏览代码

add logs for vdb-migrate command (#2626)

tags/0.5.8
Bowen Liang 1年前
父节点
当前提交
839a6a2c8a
没有帐户链接到提交者的电子邮件
共有 1 个文件被更改,包括 26 次插入3 次删除
  1. 26
    3
      api/commands.py

+ 26
- 3
api/commands.py 查看文件

""" """
click.echo(click.style('Start migrate vector db.', fg='green')) click.echo(click.style('Start migrate vector db.', fg='green'))
create_count = 0 create_count = 0
skipped_count = 0
total_count = 0
config = current_app.config config = current_app.config
vector_type = config.get('VECTOR_STORE') vector_type = config.get('VECTOR_STORE')
page = 1 page = 1


page += 1 page += 1
for dataset in datasets: for dataset in datasets:
total_count = total_count + 1
click.echo(f'Processing the {total_count} dataset {dataset.id}. '
+ f'{create_count} created, ${skipped_count} skipped.')
try: try:
click.echo('Create dataset vdb index: {}'.format(dataset.id)) click.echo('Create dataset vdb index: {}'.format(dataset.id))
if dataset.index_struct_dict: if dataset.index_struct_dict:
if dataset.index_struct_dict['type'] == vector_type: if dataset.index_struct_dict['type'] == vector_type:
skipped_count = skipped_count + 1
continue continue
collection_name = ''
if vector_type == "weaviate": if vector_type == "weaviate":
dataset_id = dataset.id dataset_id = dataset.id
collection_name = Dataset.gen_collection_name_by_id(dataset_id) collection_name = Dataset.gen_collection_name_by_id(dataset_id)
raise ValueError(f"Vector store {config.get('VECTOR_STORE')} is not supported.") raise ValueError(f"Vector store {config.get('VECTOR_STORE')} is not supported.")


vector = Vector(dataset) vector = Vector(dataset)
click.echo(f"vdb_migrate {dataset.id}")
click.echo(f"Start to migrate dataset {dataset.id}.")


try: try:
vector.delete() vector.delete()
click.echo(
click.style(f'Successfully delete vector index {collection_name} for dataset {dataset.id}.',
fg='green'))
except Exception as e: except Exception as e:
click.echo(
click.style(f'Failed to delete vector index {collection_name} for dataset {dataset.id}.',
fg='red'))
raise e raise e


dataset_documents = db.session.query(DatasetDocument).filter( dataset_documents = db.session.query(DatasetDocument).filter(
).all() ).all()


documents = [] documents = []
segments_count = 0
for dataset_document in dataset_documents: for dataset_document in dataset_documents:
segments = db.session.query(DocumentSegment).filter( segments = db.session.query(DocumentSegment).filter(
DocumentSegment.document_id == dataset_document.id, DocumentSegment.document_id == dataset_document.id,
) )


documents.append(document) documents.append(document)
segments_count = segments_count + 1


if documents: if documents:
try: try:
click.echo(click.style(
f'Start to created vector index with {len(documents)} documents of {segments_count} segments for dataset {dataset.id}.',
fg='green'))
vector.create(documents) vector.create(documents)
click.echo(
click.style(f'Successfully created vector index for dataset {dataset.id}.', fg='green'))
except Exception as e: except Exception as e:
click.echo(click.style(f'Failed to created vector index for dataset {dataset.id}.', fg='red'))
raise e raise e
click.echo(f"Dataset {dataset.id} create successfully.")
db.session.add(dataset) db.session.add(dataset)
db.session.commit() db.session.commit()
click.echo(f'Successfully migrated dataset {dataset.id}.')
create_count += 1 create_count += 1
except Exception as e: except Exception as e:
db.session.rollback() db.session.rollback()
fg='red')) fg='red'))
continue continue


click.echo(click.style('Congratulations! Create {} dataset indexes.'.format(create_count), fg='green'))
click.echo(
click.style(f'Congratulations! Create {create_count} dataset indexes, and skipped {skipped_count} datasets.',
fg='green'))




def register_commands(app): def register_commands(app):

正在加载...
取消
保存