|
@@ -131,6 +131,8 @@ def vdb_migrate():
|
|
|
"""
|
|
|
click.echo(click.style('Start migrate vector db.', fg='green'))
|
|
|
create_count = 0
|
|
|
+ skipped_count = 0
|
|
|
+ total_count = 0
|
|
|
config = current_app.config
|
|
|
vector_type = config.get('VECTOR_STORE')
|
|
|
page = 1
|
|
@@ -143,11 +145,16 @@ def vdb_migrate():
|
|
|
|
|
|
page += 1
|
|
|
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:
|
|
|
click.echo('Create dataset vdb index: {}'.format(dataset.id))
|
|
|
if dataset.index_struct_dict:
|
|
|
if dataset.index_struct_dict['type'] == vector_type:
|
|
|
+ skipped_count = skipped_count + 1
|
|
|
continue
|
|
|
+ collection_name = ''
|
|
|
if vector_type == "weaviate":
|
|
|
dataset_id = dataset.id
|
|
|
collection_name = Dataset.gen_collection_name_by_id(dataset_id)
|
|
@@ -186,11 +193,17 @@ def vdb_migrate():
|
|
|
raise ValueError(f"Vector store {config.get('VECTOR_STORE')} is not supported.")
|
|
|
|
|
|
vector = Vector(dataset)
|
|
|
- click.echo(f"vdb_migrate {dataset.id}")
|
|
|
+ click.echo(f"Start to migrate dataset {dataset.id}.")
|
|
|
|
|
|
try:
|
|
|
vector.delete()
|
|
|
+ click.echo(
|
|
|
+ click.style(f'Successfully delete vector index {collection_name} for dataset {dataset.id}.',
|
|
|
+ fg='green'))
|
|
|
except Exception as e:
|
|
|
+ click.echo(
|
|
|
+ click.style(f'Failed to delete vector index {collection_name} for dataset {dataset.id}.',
|
|
|
+ fg='red'))
|
|
|
raise e
|
|
|
|
|
|
dataset_documents = db.session.query(DatasetDocument).filter(
|
|
@@ -201,6 +214,7 @@ def vdb_migrate():
|
|
|
).all()
|
|
|
|
|
|
documents = []
|
|
|
+ segments_count = 0
|
|
|
for dataset_document in dataset_documents:
|
|
|
segments = db.session.query(DocumentSegment).filter(
|
|
|
DocumentSegment.document_id == dataset_document.id,
|
|
@@ -220,15 +234,22 @@ def vdb_migrate():
|
|
|
)
|
|
|
|
|
|
documents.append(document)
|
|
|
+ segments_count = segments_count + 1
|
|
|
|
|
|
if documents:
|
|
|
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)
|
|
|
+ click.echo(
|
|
|
+ click.style(f'Successfully created vector index for dataset {dataset.id}.', fg='green'))
|
|
|
except Exception as e:
|
|
|
+ click.echo(click.style(f'Failed to created vector index for dataset {dataset.id}.', fg='red'))
|
|
|
raise e
|
|
|
- click.echo(f"Dataset {dataset.id} create successfully.")
|
|
|
db.session.add(dataset)
|
|
|
db.session.commit()
|
|
|
+ click.echo(f'Successfully migrated dataset {dataset.id}.')
|
|
|
create_count += 1
|
|
|
except Exception as e:
|
|
|
db.session.rollback()
|
|
@@ -237,7 +258,9 @@ def vdb_migrate():
|
|
|
fg='red'))
|
|
|
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):
|