From a5c79ba9e598d1dea46e4ad82d8d6dc9a1cd3ed2 Mon Sep 17 00:00:00 2001
From: Filip Haltmayer <81822489+filip-halt@users.noreply.github.com>
Date: Tue, 9 May 2023 18:49:15 -0700
Subject: [PATCH] last few fixes (#1408)

Signed-off-by: Filip Haltmayer <filip.haltmayer@zilliz.com>
---
 pymilvus/__init__.py                    |  6 +++++-
 pymilvus/milvus_client/milvus_client.py | 21 +++++++++++++--------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/pymilvus/__init__.py b/pymilvus/__init__.py
index 545665e3b..844a6fbd2 100644
--- a/pymilvus/__init__.py
+++ b/pymilvus/__init__.py
@@ -73,6 +73,8 @@
 from .orm.future import SearchFuture, MutationFuture
 from .orm.role import Role
 
+from .milvus_client.milvus_client import MilvusClient
+
 __all__ = [
     'Collection', 'Index', 'Partition',
     'connections',
@@ -91,5 +93,7 @@
 
     'Milvus', 'Prepare', 'Status', 'DataType',
     'MilvusException',
-    '__version__'
+    '__version__',
+
+    'MilvusClient'
 ]
diff --git a/pymilvus/milvus_client/milvus_client.py b/pymilvus/milvus_client/milvus_client.py
index 0b7de6fa6..2dbe9ac1c 100644
--- a/pymilvus/milvus_client/milvus_client.py
+++ b/pymilvus/milvus_client/milvus_client.py
@@ -4,8 +4,6 @@
 from typing import Dict, List, Union
 from uuid import uuid4
 
-from tqdm import tqdm
-
 from pymilvus.exceptions import MilvusException
 from pymilvus.milvus_client.defaults import DEFAULT_SEARCH_PARAMS
 from pymilvus.orm import utility
@@ -20,17 +18,17 @@
 class MilvusClient:
     """The Milvus Client"""
 
-    # pylint: disable=logging-too-many-args, too-many-instance-attributes
+    # pylint: disable=logging-too-many-args, too-many-instance-attributes, import-outside-toplevel
 
     def __init__(
         self,
         collection_name: str = "ClientCollection",
         pk_field: str = None,
         vector_field: str = None,
-        uri: str = None,
+        uri: str = "http://localhost:19530",
         shard_num: int = None,
         partitions: List[str] = None,
-        consistency_level: str = "Bounded",
+        consistency_level: str = "Session",
         replica_number: int = 1,
         index_params: dict = None,
         timeout: int = None,
@@ -70,6 +68,14 @@ def __init__(
             overwrite (bool, optional): Whether to overwrite existing collection if exists.
                 Defaults to False
         """
+        # Optionial TQDM import
+        try:
+            import tqdm
+            self.tqdm = tqdm.tqdm
+        except ImportError:
+            logger.debug("tqdm not found")
+            self.tqdm = (lambda x, disable: x)
+
         self.uri = uri
         self.collection_name = collection_name
         self.shard_num = shard_num
@@ -166,8 +172,7 @@ def insert_data(
                 if key in self.fields:
                     insert_dict.setdefault(key, []).append(value)
 
-        # Insert the data in batches
-        for i in tqdm(range(0, len(data), batch_size), disable=not progress_bar):
+        for i in self.tqdm(range(0, len(data), batch_size), disable=not progress_bar):
             # Convert dict to list of lists batch for insertion
             try:
                 insert_batch = [
@@ -379,7 +384,7 @@ def get_vectors_by_pk(
         self,
         pks: Union[list, str, int],
         timeout: int = None,
-    ) -> None:
+    ) -> List[List[float]]:
         """Grab the inserted vectors using the primary key from the Collection.
 
         Due to current implementations, grabbing a large amount of vectors is slow.