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