feat: suppot search with highlighter in pymilvus (milvus-io#3110)

relate: milvus-io/milvus#42589

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

Author: aoiasd

pymilvus/__init__.py

@@ -45,7 +45,14 @@
 from .orm.index import Index
 from .orm.partition import Partition
 from .orm.role import Role
-from .orm.schema import CollectionSchema, FieldSchema, Function, FunctionScore, StructFieldSchema
+from .orm.schema import (
+    CollectionSchema,
+    FieldSchema,
+    Function,
+    FunctionScore,
+    LexicalHighlighter,
+    StructFieldSchema,
+)
 from .orm.utility import (
     create_resource_group,
     create_user,
@@ -96,6 +103,7 @@
     "Hits",
     "Index",
     "IndexType",
+    "LexicalHighlighter",
     "MilvusClient",
     "MilvusException",
     "MilvusUnavailableException",

pymilvus/client/async_grpc_handler.py

@@ -23,7 +23,7 @@
 )
 from pymilvus.grpc_gen import common_pb2, milvus_pb2_grpc
 from pymilvus.grpc_gen import milvus_pb2 as milvus_types
-from pymilvus.orm.schema import Function
+from pymilvus.orm.schema import Function, Highlighter
 from pymilvus.settings import Config
 
 from . import entity_helper, ts_utils, utils
@@ -828,6 +828,7 @@ async def search(
         round_decimal: int = -1,
         timeout: Optional[float] = None,
         ranker: Optional[Function] = None,
+        highlighter: Optional[Highlighter] = None,
         **kwargs,
     ):
         await self.ensure_channel_ready()
@@ -852,6 +853,7 @@ async def search(
             output_fields,
             round_decimal,
             ranker=ranker,
+            highlighter=highlighter,
             **kwargs,
         )
         return await self._execute_search(request, timeout, round_decimal=round_decimal, **kwargs)

pymilvus/client/grpc_handler.py

@@ -23,7 +23,7 @@
 )
 from pymilvus.grpc_gen import common_pb2, milvus_pb2_grpc
 from pymilvus.grpc_gen import milvus_pb2 as milvus_types
-from pymilvus.orm.schema import Function, FunctionScore
+from pymilvus.orm.schema import Function, FunctionScore, Highlighter
 from pymilvus.settings import Config
 
 from . import entity_helper, interceptor, ts_utils, utils
@@ -1010,6 +1010,7 @@ def search(
         round_decimal: int = -1,
         timeout: Optional[float] = None,
         ranker: Union[Function, FunctionScore] = None,
+        highlighter: Optional[Highlighter] = None,
         **kwargs,
     ):
         check_pass_param(
@@ -1034,6 +1035,7 @@ def search(
             output_fields,
             round_decimal,
             ranker=ranker,
+            highlighter=highlighter,
             **kwargs,
         )
         return self._execute_search(request, timeout, round_decimal=round_decimal, **kwargs)

pymilvus/client/prepare.py

@@ -16,6 +16,7 @@
     FieldSchema,
     Function,
     FunctionScore,
+    Highlighter,
     isVectorDataType,
 )
 from pymilvus.orm.types import infer_dtype_by_scalar_data
@@ -1370,6 +1371,7 @@ def search_requests_with_expr(
         output_fields: Optional[List[str]] = None,
         round_decimal: int = -1,
         ranker: Optional[Union[Function, FunctionScore]] = None,
+        highlighter: Optional[Highlighter] = None,
         **kwargs,
     ) -> milvus_types.SearchRequest:
         use_default_consistency = ts_utils.construct_guarantee_ts(collection_name, kwargs)
@@ -1524,6 +1526,9 @@ def search_requests_with_expr(
         elif ranker is not None:
             raise ParamError(message="The search ranker must be a Function or FunctionScore.")
 
+        if highlighter is not None:
+            request.highlighter.CopyFrom(Prepare.highlighter_schema(highlighter))
+
         return request
 
     @classmethod
@@ -1611,6 +1616,16 @@ def common_kv_value(v: Any) -> str:
             return json.dumps(v)
         return str(v)
 
+    @staticmethod
+    def highlighter_schema(highlighter: Highlighter) -> common_types.Highlighter:
+        return common_types.Highlighter(
+            type=highlighter.type,
+            params=[
+                common_types.KeyValuePair(key=str(k), value=Prepare.common_kv_value(v))
+                for k, v in highlighter.params.items()
+            ],
+        )
+
     @staticmethod
     def function_score_schema(function_score: FunctionScore) -> schema_types.FunctionScore:
         functions = [

pymilvus/client/search_result.py

@@ -30,6 +30,7 @@ def __init__(
         all_scores: List[float],
         fields_data: List[schema_pb2.FieldData],
         output_fields: List[str],
+        highlight_results: List[common_pb2.HighlightResult],
         pk_name: str,
     ):
         self.ids = all_pks[start:end]
@@ -93,6 +94,14 @@ def __init__(
             else:
                 msg = f"Unsupported field type: {field_data.type}"
                 raise MilvusException(msg)
+
+        if len(highlight_results) > 0:
+            for i, hit in enumerate(top_k_res):
+                hit["highlight"] = {
+                    result.field_name: list(result.datas[i + start].fragments)
+                    for result in highlight_results
+                }
+
         super().__init__(top_k_res)
 
     def __str__(self) -> str:
@@ -331,9 +340,11 @@ def _parse_search_result_data(
                     all_scores,
                     res.fields_data,
                     res.output_fields,
+                    res.highlight_results,
                     _pk_name,
                 )
             )
+
             nq_thres += topk
         return data
 
@@ -723,6 +734,10 @@ def score(self) -> float:
         """Alias of distance, will be deprecated soon"""
         return self.distance
 
+    @property
+    def highlight(self) -> Dict[str, Any]:
+        return self.data.get("highlight")
+
     @property
     def fields(self) -> Dict[str, Any]:
         """Patch for orm, will be deprecated soon"""

pymilvus/client/types.py

@@ -144,6 +144,11 @@ class FunctionType(IntEnum):
     RERANK = 3
 
 
+class HighlightType(IntEnum):
+    LEXICAL = 0
+    SEMANTIC = 1
+
+
 class RangeType(IntEnum):
     LT = 0  # less than
     LTE = 1  # less than or equal

pymilvus/grpc_gen/common_pb2.py

@@ -171,6 +171,7 @@ class MsgType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     FlushSegment: _ClassVar[MsgType]
     CreateSegment: _ClassVar[MsgType]
     Import: _ClassVar[MsgType]
+    FlushAll: _ClassVar[MsgType]
     Search: _ClassVar[MsgType]
     SearchResult: _ClassVar[MsgType]
     GetIndexState: _ClassVar[MsgType]
@@ -241,6 +242,7 @@ class MsgType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     AlterDatabase: _ClassVar[MsgType]
     DescribeDatabase: _ClassVar[MsgType]
     AddCollectionField: _ClassVar[MsgType]
+    AlterWAL: _ClassVar[MsgType]
 
 class DslType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     __slots__ = ()
@@ -357,6 +359,7 @@ class ObjectPrivilege(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     PrivilegeAddCollectionFunction: _ClassVar[ObjectPrivilege]
     PrivilegeAlterCollectionFunction: _ClassVar[ObjectPrivilege]
     PrivilegeDropCollectionFunction: _ClassVar[ObjectPrivilege]
+    PrivilegeUpdateReplicateConfiguration: _ClassVar[ObjectPrivilege]
 
 class StateCode(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     __slots__ = ()
@@ -538,6 +541,7 @@ ManualFlush: MsgType
 FlushSegment: MsgType
 CreateSegment: MsgType
 Import: MsgType
+FlushAll: MsgType
 Search: MsgType
 SearchResult: MsgType
 GetIndexState: MsgType
@@ -608,6 +612,7 @@ ListDatabases: MsgType
 AlterDatabase: MsgType
 DescribeDatabase: MsgType
 AddCollectionField: MsgType
+AlterWAL: MsgType
 Dsl: DslType
 BoolExprV1: DslType
 UndefiedState: CompactionState
@@ -706,6 +711,7 @@ PrivilegeListFileResources: ObjectPrivilege
 PrivilegeAddCollectionFunction: ObjectPrivilege
 PrivilegeAlterCollectionFunction: ObjectPrivilege
 PrivilegeDropCollectionFunction: ObjectPrivilege
+PrivilegeUpdateReplicateConfiguration: ObjectPrivilege
 Initializing: StateCode
 Healthy: StateCode
 Abnormal: StateCode