milvus-client
copied
3 changed files with 49 additions and 0 deletions
@ -0,0 +1,4 @@ |
|||
from .milvus_client import MilvusClient |
|||
|
|||
def milvus_client(*args, **kwargs): |
|||
return MilvusClient(*args, **kwargs) |
@ -0,0 +1,44 @@ |
|||
import uuid |
|||
import logging |
|||
from towhee.operator import PyOperator |
|||
from pymilvus import connections, Collection |
|||
|
|||
|
|||
logger = logging.getLogger() |
|||
|
|||
|
|||
class MilvusClient(PyOperator): |
|||
""" |
|||
Milvus ANN index class. |
|||
""" |
|||
|
|||
def __init__(self, host: str, port: int, collection_name: str): |
|||
""" |
|||
Get an existing collection. |
|||
""" |
|||
self._host = host |
|||
self._port = port |
|||
self._collection_name = collection_name |
|||
self._connect_name = uuid.uuid4().hex |
|||
connections.connect(alias=self._connect_name, host=self._host, port=self._port) |
|||
self._collection = Collection(self._collection_name, using=self._connect_name) |
|||
|
|||
def __call__(self, *data): |
|||
""" |
|||
Insert data to Milvus. |
|||
|
|||
Args: |
|||
data (`list`): |
|||
The data to insert into milvus. |
|||
|
|||
Returns: |
|||
A MutationResult object contains `insert_count` represents how many and a `primary_keys` of primary keys. |
|||
|
|||
""" |
|||
vectors = [] |
|||
for v in data: |
|||
vectors.append(v if isinstance(v, list) else [v]) |
|||
mr = self._collection.insert(vectors) |
|||
if mr.err_count > 0: |
|||
raise RuntimeError("Insert to milvus failed") |
|||
return None |
@ -0,0 +1 @@ |
|||
pymilvus |
Loading…
Reference in new issue