Unverified Commit 839a6a2c authored by Bowen Liang's avatar Bowen Liang Committed by GitHub

add logs for vdb-migrate command (#2626)

parent a769edbc
...@@ -131,6 +131,8 @@ def vdb_migrate(): ...@@ -131,6 +131,8 @@ def vdb_migrate():
""" """
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
...@@ -143,11 +145,16 @@ def vdb_migrate(): ...@@ -143,11 +145,16 @@ def vdb_migrate():
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)
...@@ -186,11 +193,17 @@ def vdb_migrate(): ...@@ -186,11 +193,17 @@ def vdb_migrate():
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(
...@@ -201,6 +214,7 @@ def vdb_migrate(): ...@@ -201,6 +214,7 @@ def vdb_migrate():
).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,
...@@ -220,15 +234,22 @@ def vdb_migrate(): ...@@ -220,15 +234,22 @@ def vdb_migrate():
) )
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()
...@@ -237,7 +258,9 @@ def vdb_migrate(): ...@@ -237,7 +258,9 @@ def vdb_migrate():
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):
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment