diff --git a/pymilvus/orm/connections.py b/pymilvus/orm/connections.py index af71d8df9..761ae2abc 100644 --- a/pymilvus/orm/connections.py +++ b/pymilvus/orm/connections.py @@ -12,6 +12,7 @@ import copy import logging +import pathlib import threading import time from typing import Callable, Tuple, Union @@ -357,6 +358,32 @@ 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: {parent_path} not exists" + ) + + from milvus_lite.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 diff --git a/pyproject.toml b/pyproject.toml index 9bfaa3f16..08d8858eb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ dependencies=[ "minio>=7.0.0", "pyarrow>=12.0.0", "azure-storage-blob", + "milvus-lite>=2.4.0", ] classifiers=[ diff --git a/requirements.txt b/requirements.txt index 85a1c20b3..d5251ad49 100644 --- a/requirements.txt +++ b/requirements.txt @@ -37,3 +37,4 @@ black requests minio azure-storage-blob +milvus-lite>=2.4.0 diff --git a/tests/test_milvus_lite.py b/tests/test_milvus_lite.py new file mode 100644 index 000000000..ef4f49274 --- /dev/null +++ b/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