From f25927855ea34294ffe13c7772849b4026d1902e Mon Sep 17 00:00:00 2001
From: Jyong <76649700+JohnJyong@users.noreply.github.com>
Date: Wed, 22 May 2024 01:42:08 +0800
Subject: [PATCH] add qdrant metadata.doc_id index (#4559)

---
 api/commands.py | 55 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 55 insertions(+)

diff --git a/api/commands.py b/api/commands.py
index 75f2491421..86cc9e204f 100644
--- a/api/commands.py
+++ b/api/commands.py
@@ -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)
+