update

langgenius · Jul 15, 2024 · 997901e · 997901e
1 parent 4fdcb30
commit 997901e
Show file tree

Hide file tree

Showing 4 changed files with 84 additions and 132 deletions.
diff --git a/api/core/rag/datasource/vdb/milvus/milvus_vector.py b/api/core/rag/datasource/vdb/milvus/milvus_vector.py
@@ -6,6 +6,7 @@
 from flask import current_app
 from pydantic import BaseModel, model_validator
 from pymilvus import MilvusClient, MilvusException, connections
+from pymilvus.milvus_client import IndexParams
 
 from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.field import Field
@@ -248,11 +249,15 @@ def create_collection(
                 # Since primary field is auto-id, no need to track it
                 self._fields.remove(Field.PRIMARY_KEY.value)
 
+                # Create Index params for the collection
+                index_params_obj = IndexParams()
+                index_params_obj.add_index(field_name=Field.VECTOR.value, **index_params)
+
                 # Create the collection
                 collection_name = self._collection_name
-                self._client.create_collection_with_schema(collection_name=collection_name,
-                                                           schema=schema, index_param=index_params,
-                                                           consistency_level=self._consistency_level)
+                self._client.create_collection(collection_name=collection_name,
+                                               schema=schema, index_params=index_params_obj,
+                                               consistency_level=self._consistency_level)
             redis_client.set(collection_exist_cache_key, 1, ex=3600)
 
     def _init_client(self, config) -> MilvusClient: