Support milvuslite

Signed-off-by: junjie.jiang <[email protected]>

pymilvus/orm/connections.py

@@ -15,6 +15,7 @@
 import threading
 import time
 from typing import Callable, Tuple, Union
+import pathlib
 from urllib import parse
 
 from pymilvus.client.check import is_legal_address, is_legal_host, is_legal_port
@@ -357,6 +358,20 @@ def connect(
             >>> connections.connect("test", host="localhost", port="19530")
         """
 
+        if kwargs.get("uri") and parse.urlparse(kwargs["uri"]).scheme.lower() not in ["unix", "http", "https", "tcp"]:
+            # start and connect milvuslite
+            if kwargs["uri"].endswith("/"):
+                raise ConnectionConfigException(message=f"Open local milvus failed, {kwargs['uri']} is not a local file path")
+            parent_path = pathlib.Path(kwargs["uri"]).parent
+            if not parent_path.is_dir():
+                raise ConnectionConfigException(message=f"Open local milvus failed, dir: {str(parent_path)} not exists")
+
+            from milvus.server_manager import server_manager_instance
+            local_uri = server_manager_instance.start_and_get_uri(kwargs["uri"])
+            if local_uri is None:
+                raise ConnectionConfigException(message="Open local milvus failed")
+            kwargs["uri"] = local_uri
+
         # kwargs_copy is used for auto reconnect
         kwargs_copy = copy.deepcopy(kwargs)
         kwargs_copy["user"] = user

pyproject.toml

@@ -29,6 +29,7 @@ dependencies=[
     "pyarrow>=12.0.0",
     "azure-storage-blob",
     "scipy",
+    "milvus>=2.4.0,<2.5.0",
 ]
 
 classifiers=[

requirements.txt

@@ -38,3 +38,4 @@ black
 requests
 minio
 azure-storage-blob
+milvus>=2.4.0,<2.5.0

tests/test_milvus_lite.py

@@ -0,0 +1,56 @@
+import os
+from tempfile import TemporaryDirectory
+import numpy as np
+
+from pymilvus.milvus_client import MilvusClient
+
+
+class TestMilvusLite:
+
+    def test_milvus_lite(self):
+        with TemporaryDirectory(dir='./') as root:
+            db_file = os.path.join(root, 'test.db')
+            client = MilvusClient(db_file)
+            client.create_collection(
+                collection_name="demo_collection",
+                dimension=3
+            )
+
+            # Text strings to search from.
+            docs = [
+                "Artificial intelligence was founded as an academic discipline in 1956.",
+                "Alan Turing was the first person to conduct substantial research in AI.",
+                "Born in Maida Vale, London, Turing was raised in southern England.",
+            ]
+
+            vectors = [[np.random.uniform(-1, 1) for _ in range(3) ] for _ in range(len(docs))]
+            data = [{"id": i, "vector": vectors[i], "text": docs[i], "subject": "history"} for i in range(len(vectors))]
+            res = client.insert(
+                collection_name="demo_collection",
+                data=data
+            )
+            assert res["insert_count"] == 3
+
+            res = client.search(
+                    collection_name="demo_collection",
+                    data=[vectors[0]],
+                    filter="subject == 'history'",
+                    limit=2,
+                    output_fields=["text", "subject"],
+                )
+            assert len(res[0]) == 2
+
+            # a query that retrieves all entities matching filter expressions.
+            res = client.query(
+                    collection_name="demo_collection",
+                    filter="subject == 'history'",
+                    output_fields=["text", "subject"],
+                )
+            assert len(res) == 3
+
+            # delete
+            res = client.delete(
+                    collection_name="demo_collection",
+                    filter="subject == 'history'",
+                )
+            assert len(res) == 3