Skip to content

Commit

Permalink
Better naming and creds access
Browse files Browse the repository at this point in the history
Signed-off-by: Filip Haltmayer <[email protected]>
  • Loading branch information
Filip Haltmayer committed Jun 27, 2023
1 parent 6654dc1 commit bc72755
Showing 1 changed file with 36 additions and 13 deletions.
49 changes: 36 additions & 13 deletions milvus.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@
SimilarityIndex,
)

import fiftyone.brain.internal.core.utils as fbu


pymilvus = fou.lazy_import("pymilvus")


Expand All @@ -28,6 +31,7 @@
"euclidean": "L2",
}


class MilvusSimilarityConfig(SimilarityConfig):
"""Configuration for the Milvus similarity backend.
Expand All @@ -47,20 +51,22 @@ class MilvusSimilarityConfig(SimilarityConfig):
uri (str): full address of Milvus server.
user (str): username if using rbac.
password(str): password for supplied username.
consistency_level(str): which consistency level to use. Possible values are Strong, Session, Bounded, Eventually.
consistency_level(str): which consistency level to use. Possible values are
Strong, Session, Bounded, Eventually.
overwrite(str): whether to overwrite the collection if it already exists.
"""

def __init__(
self,
embeddings_field=None,
model=None,
patches_field=None,
supports_prompts=None,
metric="euclidean",
collection_name: str = "ClientCollection",
collection_name: str = None,
uri: str = "http://localhost:19530",
user: str = "",
password: str = "",
user: str = None,
password: str = None,
consistency_level: str = "Session",
**kwargs,
):
Expand All @@ -80,8 +86,8 @@ def __init__(
self.metric = metric
self.collection_name = collection_name
self._uri = uri
self.user = user
self.password = password
self._user = user
self._password = password
self.consistency_level = consistency_level
self.index_params = {
"metric_type": _SUPPORTED_METRICS[metric],
Expand All @@ -104,6 +110,22 @@ def uri(self):
def uri(self, uri):
self._uri = uri

@property
def user(self):
return self._user

@user.setter
def user(self, user):
self._user = user

@property
def password(self):
return self._password

@password.setter
def password(self, password):
self._password = password

@property
def max_k(self):
return 16_384
Expand Down Expand Up @@ -158,6 +180,14 @@ def _initialize(self):
self.config.uri, self.config.user, self.config.password
)

if self.config.collection_name is None:
root = "fiftyone-" + fou.to_slug(self.samples._root_dataset.name)
collection_name = fbu.get_unique_name(
root, utility.list_collections(using=self.alias)
)
self.config.collection_name = collection_name.replace("-", "_")
self.save_config()

self._init_collection()

def _connect(self, uri, user, password):
Expand All @@ -179,13 +209,6 @@ def _init_collection(self):
if utility.has_collection(self.config.collection_name, using=self.alias):
col = Collection(self.config.collection_name, using=self.alias)
col.load()
for x in col.schema.fields:
if x.params.get("dim", None) is not None:
dim = x.params["dim"]
break
return (col, dim)

return None, None

@property
def total_index_size(self):
Expand Down

0 comments on commit bc72755

Please sign in to comment.