enhance: support run analyzer (milvus-io#2622)

relate: milvus-io/milvus#39705

Signed-off-by: aoiasd <[email protected]>

Author: aoiasd

pymilvus/client/grpc_handler.py

@@ -45,6 +45,7 @@
 from .prepare import Prepare
 from .search_reasult import SearchResult
 from .types import (
+    AnalyzeResult,
     BulkInsertState,
     CompactionPlans,
     CompactionState,
@@ -2243,3 +2244,20 @@ def remove_privileges_from_group(
         )
         resp = self._stub.OperatePrivilegeGroup(req, wait_for_ready=True, timeout=timeout)
         check_status(resp)
+
+    @retry_on_rpc_failure()
+    def run_analyzer(
+        self,
+        texts: Union[str, List[str]],
+        analyzer_params: Union[str, Dict],
+        timeout: Optional[float] = None,
+        **kwargs,
+    ):
+        check_pass_param(timeout=timeout)
+        req = Prepare.run_analyzer(texts, analyzer_params)
+        resp = self._stub.RunAnalyzer(req, timeout=timeout)
+        check_status(resp.status)
+
+        if isinstance(texts, str):
+            return AnalyzeResult(resp.results[0])
+        return [AnalyzeResult(result) for result in resp.results]

pymilvus/client/prepare.py

@@ -1724,3 +1724,18 @@ def operate_privilege_group_req(
             privileges=[milvus_types.PrivilegeEntity(name=p) for p in privileges],
             type=operate_privilege_group_type,
         )
+
+    @classmethod
+    def run_analyzer(cls, texts: Union[str, List[str]], analyzer_params: Union[str, Dict]):
+        req = milvus_types.RunAnalyzerRequset()
+        if isinstance(texts, str):
+            req.placeholder.append(texts.encode("utf-8"))
+        else:
+            req.placeholder.extend([text.encode("utf-8") for text in texts])
+
+        if isinstance(analyzer_params, dict):
+            req.analyzer_params = ujson.dumps(analyzer_params)
+        else:
+            req.analyzer_params = analyzer_params
+
+        return req

pymilvus/client/types.py

@@ -1081,3 +1081,13 @@ def to_dict(self) -> Dict[str, Any]:
         result = {"name": self.name}
         result.update(self.properties)
         return result
+
+
+class AnalyzeResult:
+    def __init__(self, info: milvus_types.AnalyzerResult) -> None:
+        self.tokens = info.tokens
+
+    def __str__(self) -> str:
+        return f"{{tokens: {self.tokens}}}"
+
+    __reper__ = __str__

pymilvus/grpc_gen/milvus-proto

@@ -2219,3 +2219,47 @@ class ListImportsAuthPlaceholder(_message.Message):
     db_name: str
     collection_name: str
     def __init__(self, db_name: _Optional[str] = ..., collection_name: _Optional[str] = ...) -> None: ...
+
+class RunAnalyzerRequest(_message.Message):
+    __slots__ = ("base", "analyzer_params", "placeholder", "with_detail", "with_hash")
+    BASE_FIELD_NUMBER: _ClassVar[int]
+    ANALYZER_PARAMS_FIELD_NUMBER: _ClassVar[int]
+    PLACEHOLDER_FIELD_NUMBER: _ClassVar[int]
+    WITH_DETAIL_FIELD_NUMBER: _ClassVar[int]
+    WITH_HASH_FIELD_NUMBER: _ClassVar[int]
+    base: _common_pb2.MsgBase
+    analyzer_params: str
+    placeholder: _containers.RepeatedScalarFieldContainer[bytes]
+    with_detail: bool
+    with_hash: bool
+    def __init__(self, base: _Optional[_Union[_common_pb2.MsgBase, _Mapping]] = ..., analyzer_params: _Optional[str] = ..., placeholder: _Optional[_Iterable[bytes]] = ..., with_detail: bool = ..., with_hash: bool = ...) -> None: ...
+
+class AnalyzerToken(_message.Message):
+    __slots__ = ("token", "start_offset", "end_offset", "position", "position_length", "hash")
+    TOKEN_FIELD_NUMBER: _ClassVar[int]
+    START_OFFSET_FIELD_NUMBER: _ClassVar[int]
+    END_OFFSET_FIELD_NUMBER: _ClassVar[int]
+    POSITION_FIELD_NUMBER: _ClassVar[int]
+    POSITION_LENGTH_FIELD_NUMBER: _ClassVar[int]
+    HASH_FIELD_NUMBER: _ClassVar[int]
+    token: str
+    start_offset: int
+    end_offset: int
+    position: int
+    position_length: int
+    hash: int
+    def __init__(self, token: _Optional[str] = ..., start_offset: _Optional[int] = ..., end_offset: _Optional[int] = ..., position: _Optional[int] = ..., position_length: _Optional[int] = ..., hash: _Optional[int] = ...) -> None: ...
+
+class AnalyzerResult(_message.Message):
+    __slots__ = ("tokens",)
+    TOKENS_FIELD_NUMBER: _ClassVar[int]
+    tokens: _containers.RepeatedCompositeFieldContainer[AnalyzerToken]
+    def __init__(self, tokens: _Optional[_Iterable[_Union[AnalyzerToken, _Mapping]]] = ...) -> None: ...
+
+class RunAnalyzerResponse(_message.Message):
+    __slots__ = ("status", "results")
+    STATUS_FIELD_NUMBER: _ClassVar[int]
+    RESULTS_FIELD_NUMBER: _ClassVar[int]
+    status: _common_pb2.Status
+    results: _containers.RepeatedCompositeFieldContainer[AnalyzerResult]
+    def __init__(self, status: _Optional[_Union[_common_pb2.Status, _Mapping]] = ..., results: _Optional[_Iterable[_Union[AnalyzerResult, _Mapping]]] = ...) -> None: ...

pymilvus/grpc_gen/milvus_pb2.py

@@ -486,6 +486,11 @@ def __init__(self, channel):
                 request_serializer=milvus__pb2.OperatePrivilegeGroupRequest.SerializeToString,
                 response_deserializer=common__pb2.Status.FromString,
                 )
+        self.RunAnalyzer = channel.unary_unary(
+                '/milvus.proto.milvus.MilvusService/RunAnalyzer',
+                request_serializer=milvus__pb2.RunAnalyzerRequest.SerializeToString,
+                response_deserializer=milvus__pb2.RunAnalyzerResponse.FromString,
+                )
 
 
 class MilvusServiceServicer(object):
@@ -1062,6 +1067,12 @@ def OperatePrivilegeGroup(self, request, context):
         context.set_details('Method not implemented!')
         raise NotImplementedError('Method not implemented!')
 
+    def RunAnalyzer(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
 
 def add_MilvusServiceServicer_to_server(servicer, server):
     rpc_method_handlers = {
@@ -1535,6 +1546,11 @@ def add_MilvusServiceServicer_to_server(servicer, server):
                     request_deserializer=milvus__pb2.OperatePrivilegeGroupRequest.FromString,
                     response_serializer=common__pb2.Status.SerializeToString,
             ),
+            'RunAnalyzer': grpc.unary_unary_rpc_method_handler(
+                    servicer.RunAnalyzer,
+                    request_deserializer=milvus__pb2.RunAnalyzerRequest.FromString,
+                    response_serializer=milvus__pb2.RunAnalyzerResponse.SerializeToString,
+            ),
     }
     generic_handler = grpc.method_handlers_generic_handler(
             'milvus.proto.milvus.MilvusService', rpc_method_handlers)
@@ -3143,6 +3159,23 @@ def OperatePrivilegeGroup(request,
             options, channel_credentials,
             insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
 
+    @staticmethod
+    def RunAnalyzer(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(request, target, '/milvus.proto.milvus.MilvusService/RunAnalyzer',
+            milvus__pb2.RunAnalyzerRequest.SerializeToString,
+            milvus__pb2.RunAnalyzerResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
 
 class ProxyServiceStub(object):
     """Missing associated documentation comment in .proto file."""

pymilvus/grpc_gen/milvus_pb2.pyi

@@ -1716,3 +1716,22 @@ def describe_replica(
         """
         conn = self._get_connection()
         return conn.describe_replica(collection_name, timeout=timeout, **kwargs)
+
+    def run_analyzer(
+        self,
+        texts: Union[str, List[str]],
+        analyzer_params: Union[str, Dict, None] = None,
+        timeout: Optional[float] = None,
+    ):
+        """Run analyzer. Return result tokens of analysis.
+        Args:
+            text(``str``,``List[str]``): The input text (string or string list).
+            analyzer_params(``str``,``Dict``,``None``): The parameters of analyzer.
+            timeout(``float``, optional): The timeout value in seconds. Defaults to None.
+        Returns:
+                (``List[str]``,``List[List[str]]``): The result tokens of analysis.
+        """
+        if analyzer_params is None:
+            analyzer_params = {}
+
+        return self._get_connection().run_analyzer(texts, analyzer_params, timeout=timeout)