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