🚧 fix(wip): good milvus changes

Prismadic · Mar 15, 2024 · 8eba19d · 8eba19d
1 parent 677de74
commit 8eba19d
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 10 deletions.
diff --git a/magnet/ize/memory.py b/magnet/ize/memory.py
@@ -31,12 +31,12 @@ async def on(self, create: bool = False, initialize: bool = False):
         self._model = SentenceTransformer(self.config.index.model, device=Utils().check_cuda())
         _f('info', f'loading into {self._model.device}')
         self.db = MilvusDB(self.config)
-        self.db.on()
+        await self.db.on()
         if create:
-            self.db.create(overwrite=True)
-        self.db.load()
+            await self.db.create(overwrite=True)
+        await self.db.load()
         if initialize:
-                self.db.initialize()
+            self.db.initialize()
 
     async def index(self, payload, msg, field: Charge = None, v: bool = False, instruction: str = "Represent this sentence for searching relevant passages: "):
         if not msg or not payload:

diff --git a/magnet/utils/index/milvus.py b/magnet/utils/index/milvus.py
@@ -12,7 +12,7 @@ def __init__(self, config: MagnetConfig):
             FieldSchema(name='text', dtype=DataType.VARCHAR, max_length=65535),
             FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR, dim=self.config.index.dimension)
         ]
-    def on(self):
+    async def on(self):
         try:
             _f('wait', f'connecting to {self.config.index.milvus_uri}')
             self.connection = connections.connect(
@@ -34,17 +34,17 @@ async def off(self):
         except Exception as e:
             return _f('fatal', e)
 
-    def create(self, overwrite=False):
-        if utility.has_collection(self.config.index.name, using=self.config.session) and overwrite:
-            utility.drop_collection(self.config.index.name, using=self.config.session)
+    async def create(self, overwrite=False):
         try:
+            if utility.has_collection(self.config.index.name, using=self.config.session) and overwrite:
+                utility.drop_collection(self.config.index.name, using=self.config.session)
             self.collection = Collection(name=self.config.index.name, schema=self.schema, using=self.config.session)
             self.collection.create_index(field_name="embedding", index_params=self.config.index.options)
             _f('success', f"{self.config.index.name} created")
         except Exception as e:
             _f('fatal', e)
 
-    def load(self):
+    async def load(self):
         _f('wait', f'loading {self.config.index.name} into memory, may take time')
         self.collection = Collection(name=self.config.index.name, schema=self.schema, using=self.config.session)
         self.collection.load()
@@ -71,7 +71,7 @@ def initialize(self, user: str = 'magnet', password: str = '33011033'):
         except Exception as e:
             _f('fatal', e)
 
-    def delete_index(self):
+    async def delete_index(self):
         if utility.has_collection(self.config.index.name, using=self.config.session):
             utility.drop_collection(self.config.index.name, using=self.config.session)
             _f('warn', f"Index for {self.config.index.name} deleted")