转移

models/milvus.py

@@ -0,0 +1,26 @@
+import pymilvus
+
+from configs.config import MILVUS_HOST
+
+# 连接 Milvus (开启 Milvus 服务)
+collection_name = 'default'
+
+
+# 获取 Milvus 连接
+def get_collection(collection_name):
+    pymilvus.connections.connect(host=MILVUS_HOST, port='19530')
+    if not pymilvus.utility.has_collection(collection_name):
+        field1 = pymilvus.FieldSchema(name="id", dtype=pymilvus.DataType.INT64, is_primary=True)
+        field2 = pymilvus.FieldSchema(name="embedding", dtype=pymilvus.DataType.FLOAT_VECTOR, dim=2048)
+        field3 = pymilvus.FieldSchema(name="article_id", dtype=pymilvus.DataType.INT64)
+        schema = pymilvus.CollectionSchema(fields=[field1, field2, field3])
+        return pymilvus.Collection(name=collection_name, schema=schema)
+    return pymilvus.Collection(name=collection_name)
+
+
+# 检查索引是否存在, 不存在则创建, 并加载
+#collection = get_collection(collection_name)
+#if not collection.has_index():
+#    default_index = {"index_type": "IVF_SQ8", "metric_type": 'L2', "params": {"nlist": 16384}}
+#    collection.create_index(field_name="embedding", index_params=default_index)
+#collection.load()