|
|
@@ -448,9 +448,64 @@ def convert_to_agent_apps():
|
|
|
click.echo(click.style('Congratulations! Converted {} agent apps.'.format(len(proceeded_app_ids)), fg='green'))
|
|
|
|
|
|
|
|
|
+@click.command('add-qdrant-doc-id-index', help='add qdrant doc_id index.')
|
|
|
+@click.option('--field', default='metadata.doc_id', prompt=False, help='index field , default is metadata.doc_id.')
|
|
|
+def add_qdrant_doc_id_index(field: str):
|
|
|
+ click.echo(click.style('Start add qdrant doc_id index.', fg='green'))
|
|
|
+ config = current_app.config
|
|
|
+ vector_type = config.get('VECTOR_STORE')
|
|
|
+ if vector_type != "qdrant":
|
|
|
+ click.echo(click.style('Sorry, only support qdrant vector store.', fg='red'))
|
|
|
+ return
|
|
|
+ create_count = 0
|
|
|
+
|
|
|
+ try:
|
|
|
+ bindings = db.session.query(DatasetCollectionBinding).all()
|
|
|
+ if not bindings:
|
|
|
+ click.echo(click.style('Sorry, no dataset collection bindings found.', fg='red'))
|
|
|
+ return
|
|
|
+ import qdrant_client
|
|
|
+ from qdrant_client.http.exceptions import UnexpectedResponse
|
|
|
+ from qdrant_client.http.models import PayloadSchemaType
|
|
|
+
|
|
|
+ from core.rag.datasource.vdb.qdrant.qdrant_vector import QdrantConfig
|
|
|
+ for binding in bindings:
|
|
|
+ qdrant_config = QdrantConfig(
|
|
|
+ endpoint=config.get('QDRANT_URL'),
|
|
|
+ api_key=config.get('QDRANT_API_KEY'),
|
|
|
+ root_path=current_app.root_path,
|
|
|
+ timeout=config.get('QDRANT_CLIENT_TIMEOUT'),
|
|
|
+ grpc_port=config.get('QDRANT_GRPC_PORT'),
|
|
|
+ prefer_grpc=config.get('QDRANT_GRPC_ENABLED')
|
|
|
+ )
|
|
|
+ try:
|
|
|
+ client = qdrant_client.QdrantClient(**qdrant_config.to_qdrant_params())
|
|
|
+ # create payload index
|
|
|
+ client.create_payload_index(binding.collection_name, field,
|
|
|
+ field_schema=PayloadSchemaType.KEYWORD)
|
|
|
+ create_count += 1
|
|
|
+ except UnexpectedResponse as e:
|
|
|
+ # Collection does not exist, so return
|
|
|
+ if e.status_code == 404:
|
|
|
+ click.echo(click.style(f'Collection not found, collection_name:{binding.collection_name}.', fg='red'))
|
|
|
+ continue
|
|
|
+ # Some other error occurred, so re-raise the exception
|
|
|
+ else:
|
|
|
+ click.echo(click.style(f'Failed to create qdrant index, collection_name:{binding.collection_name}.', fg='red'))
|
|
|
+
|
|
|
+ except Exception as e:
|
|
|
+ click.echo(click.style('Failed to create qdrant client.', fg='red'))
|
|
|
+
|
|
|
+ click.echo(
|
|
|
+ click.style(f'Congratulations! Create {create_count} collection indexes.',
|
|
|
+ fg='green'))
|
|
|
+
|
|
|
+
|
|
|
def register_commands(app):
|
|
|
app.cli.add_command(reset_password)
|
|
|
app.cli.add_command(reset_email)
|
|
|
app.cli.add_command(reset_encrypt_key_pair)
|
|
|
app.cli.add_command(vdb_migrate)
|
|
|
app.cli.add_command(convert_to_agent_apps)
|
|
|
+ app.cli.add_command(add_qdrant_doc_id_index)
|
|
|
+
|
Jyong