fix(brains+milvus): fixed brain_id + milvus index creation error

Author: ChrisCoder9000

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "brainapi2"
-version = "1.6.14-dev"
+version = "1.6.16-dev"
 description = "Version 1.x.x of the BrainAPI memory layer."
 authors = [
     {name = "Christian",email = "[email protected]"}

src/lib/milvus/client.py

@@ -92,12 +92,28 @@ def _ensure_store(self, store: str, brain_id: str) -> None:
         if store not in EMBEDDING_STORES_SIZES:
             raise ValueError(f"Store {store} not available")
 
+        collection_created = False
         if not self.client.has_collection(store):
             self.client.create_collection(
                 store,
                 dimension=EMBEDDING_STORES_SIZES[store],
                 vector_field_name="embeddings",
             )
+            collection_created = True
+
+        if collection_created:
+            try:
+                index_params = {
+                    "metric_type": "COSINE",
+                    "index_type": "AUTOINDEX",
+                }
+                self.client.create_index(
+                    collection_name=store,
+                    field_name="embeddings",
+                    index_params=index_params,
+                )
+            except Exception as e:
+                pass
 
         try:
             has_partition = self.client.has_partition(store, brain_id)

src/services/api/middlewares/brains.py

@@ -82,21 +82,21 @@ async def receive():
         if brain_id and not cached_brain_id:
             stored_brain = data_adapter.get_brain(name_key=brain_id)
 
+            request.state.brain_id = brain_id
+
             if not stored_brain and brain_creation_allowed:
                 new_brain = data_adapter.create_brain(name_key=brain_id)
                 cache_adapter.set(
                     key=f"brain:{brain_id}",
                     value=new_brain.id,
                     brain_id="system",
                 )
-                request.state.brain_id = new_brain.id
             elif stored_brain:
                 cache_adapter.set(
                     key=f"brain:{brain_id}",
                     value=stored_brain.id,
                     brain_id="system",
                 )
-                request.state.brain_id = stored_brain.id
         elif not brain_id and default_brain_fallback:
             default_brain = data_adapter.get_brain(name_key="default")
             if not default_brain: