Просмотр исходного кода

add logs for vdb-migrate command (#2626)

tags/0.5.8
Bowen Liang 1 год назад
Родитель
Сommit
839a6a2c8a
Аккаунт пользователя с таким Email не найден
1 измененных файлов: 26 добавлений и 3 удалений
  1. 26
    3
      api/commands.py

+ 26
- 3
api/commands.py Просмотреть файл

@@ -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):

Загрузка…
Отмена
Сохранить