Support function edit sdk

Signed-off-by: junjie.jiang <junjie.jiang@zilliz.com>

Author: junjiejiangjjj

examples/function_edit.py

@@ -0,0 +1,62 @@
+from pymilvus import (
+    MilvusClient,
+    Function, DataType, FunctionType,
+)
+
+collection_name = "text_embedding"
+
+milvus_client = MilvusClient("http://localhost:19530")
+
+has_collection = milvus_client.has_collection(collection_name, timeout=5)
+if has_collection:
+    milvus_client.drop_collection(collection_name)
+
+schema = milvus_client.create_schema()
+schema.add_field("id", DataType.INT64, is_primary=True, auto_id=False)
+schema.add_field("document", DataType.VARCHAR, max_length=9000)
+schema.add_field("dense", DataType.FLOAT_VECTOR, dim=1536)
+
+text_embedding_function = Function(
+    name="openai",
+    function_type=FunctionType.TEXTEMBEDDING,
+    input_field_names=["document"],
+    output_field_names="dense",
+    params={
+        "provider": "openai",
+        "model_name": "text-embedding-3-small",
+    }
+)
+
+schema.add_function(text_embedding_function)
+
+index_params = milvus_client.prepare_index_params()
+index_params.add_index(
+    field_name="dense",
+    index_name="dense_index",
+    index_type="AUTOINDEX",
+    metric_type="IP",
+)
+
+ret = milvus_client.create_collection(collection_name, schema=schema, index_params=index_params, consistency_level="Strong")
+
+ret = milvus_client.describe_collection(collection_name)
+print(ret["functions"][0])
+
+text_embedding_function.params["user"] = "user123"
+
+milvus_client.alter_collection_function(collection_name, "openai", text_embedding_function)
+
+ret = milvus_client.describe_collection(collection_name)
+print(ret["functions"][0])
+
+milvus_client.drop_collection_function(collection_name, "openai")
+
+ret = milvus_client.describe_collection(collection_name)
+print(ret["functions"])
+
+text_embedding_function.params["user"] = "user1234"
+
+milvus_client.add_collection_function(collection_name, text_embedding_function)
+
+ret = milvus_client.describe_collection(collection_name)
+print(ret["functions"][0])

pymilvus/client/async_grpc_handler.py

@@ -1310,6 +1310,60 @@ async def add_collection_field(
         )
         check_status(status)
 
+    @retry_on_rpc_failure()
+    async def drop_collection_function(
+        self,
+        collection_name: str,
+        function_name: str,
+        timeout: Optional[float] = None,
+        **kwargs,
+    ):
+        await self.ensure_channel_ready()
+        check_pass_param(collection_name=collection_name, timeout=timeout)
+        request = Prepare.drop_collection_function_request(collection_name, function_name)
+
+        status = await self._async_stub.DropCollectionFunction(
+            request, timeout=timeout, metadata=_api_level_md(**kwargs)
+        )
+        check_status(status)
+
+    @retry_on_rpc_failure()
+    async def add_collection_function(
+        self,
+        collection_name: str,
+        function: Function,
+        timeout: Optional[float] = None,
+        **kwargs,
+    ):
+        await self.ensure_channel_ready()
+        check_pass_param(collection_name=collection_name, timeout=timeout)
+        request = Prepare.add_collection_function_request(collection_name, function)
+
+        status = await self._async_stub.AddCollectionFunction(
+            request, timeout=timeout, metadata=_api_level_md(**kwargs)
+        )
+        check_status(status)
+
+    @retry_on_rpc_failure()
+    async def alter_collection_function(
+        self,
+        collection_name: str,
+        function_name: str,
+        function: Function,
+        timeout: Optional[float] = None,
+        **kwargs,
+    ):
+        await self.ensure_channel_ready()
+        check_pass_param(collection_name=collection_name, timeout=timeout)
+        request = Prepare.alter_collection_function_request(
+            collection_name, function_name, function
+        )
+
+        status = await self._async_stub.AlterCollectionFunction(
+            request, timeout=timeout, metadata=_api_level_md(**kwargs)
+        )
+        check_status(status)
+
     @retry_on_rpc_failure()
     async def list_indexes(self, collection_name: str, timeout: Optional[float] = None, **kwargs):
         await self.ensure_channel_ready()

pymilvus/client/grpc_handler.py

@@ -346,6 +346,57 @@ def add_collection_field(
         )
         check_status(status)
 
+    @retry_on_rpc_failure()
+    def drop_collection_function(
+        self,
+        collection_name: str,
+        function_name: str,
+        timeout: Optional[float] = None,
+        **kwargs,
+    ):
+        check_pass_param(collection_name=collection_name, timeout=timeout)
+        request = Prepare.drop_collection_function_request(collection_name, function_name)
+
+        status = self._stub.DropCollectionFunction(
+            request, timeout=timeout, metadata=_api_level_md(**kwargs)
+        )
+        check_status(status)
+
+    @retry_on_rpc_failure()
+    def add_collection_function(
+        self,
+        collection_name: str,
+        function: Function,
+        timeout: Optional[float] = None,
+        **kwargs,
+    ):
+        check_pass_param(collection_name=collection_name, timeout=timeout)
+        request = Prepare.add_collection_function_request(collection_name, function)
+
+        status = self._stub.AddCollectionFunction(
+            request, timeout=timeout, metadata=_api_level_md(**kwargs)
+        )
+        check_status(status)
+
+    @retry_on_rpc_failure()
+    def alter_collection_function(
+        self,
+        collection_name: str,
+        function_name: str,
+        function: Function,
+        timeout: Optional[float] = None,
+        **kwargs,
+    ):
+        check_pass_param(collection_name=collection_name, timeout=timeout)
+        request = Prepare.alter_collection_function_request(
+            collection_name, function_name, function
+        )
+
+        status = self._stub.AlterCollectionFunction(
+            request, timeout=timeout, metadata=_api_level_md(**kwargs)
+        )
+        check_status(status)
+
     @retry_on_rpc_failure()
     def alter_collection_properties(
         self, collection_name: str, properties: List, timeout: Optional[float] = None, **kwargs

pymilvus/client/prepare.py

@@ -241,16 +241,7 @@ def get_schema_from_collection_schema(
             schema.struct_array_fields.append(struct_schema)
 
         for f in fields.functions:
-            function_schema = schema_types.FunctionSchema(
-                name=f.name,
-                description=f.description,
-                type=f.type,
-                input_field_names=f.input_field_names,
-                output_field_names=f.output_field_names,
-            )
-            for k, v in f.params.items():
-                kv_pair = common_types.KeyValuePair(key=str(k), value=str(v))
-                function_schema.params.append(kv_pair)
+            function_schema = cls.convert_function_to_function_schema(f)
             schema.functions.append(function_schema)
 
         return schema
@@ -363,6 +354,34 @@ def get_schema(
     def drop_collection_request(cls, collection_name: str) -> milvus_types.DropCollectionRequest:
         return milvus_types.DropCollectionRequest(collection_name=collection_name)
 
+    @classmethod
+    def drop_collection_function_request(
+        cls, collection_name: str, function_name: str
+    ) -> milvus_types.DropCollectionFunctionRequest:
+        return milvus_types.DropCollectionFunctionRequest(
+            collection_name=collection_name, function_name=function_name
+        )
+
+    @classmethod
+    def add_collection_function_request(
+        cls, collection_name: str, f: Function
+    ) -> milvus_types.AddCollectionFunctionRequest:
+        function_schema = cls.convert_function_to_function_schema(f)
+        return milvus_types.AddCollectionFunctionRequest(
+            collection_name=collection_name, functionSchema=function_schema
+        )
+
+    @classmethod
+    def alter_collection_function_request(
+        cls, collection_name: str, function_name: str, f: Function
+    ) -> milvus_types.AlterCollectionFunctionRequest:
+        function_schema = cls.convert_function_to_function_schema(f)
+        return milvus_types.AlterCollectionFunctionRequest(
+            collection_name=collection_name,
+            function_name=function_name,
+            functionSchema=function_schema,
+        )
+
     @classmethod
     def add_collection_field_request(
         cls,
@@ -2424,3 +2443,17 @@ def update_replicate_configuration_request(
         return milvus_types.UpdateReplicateConfigurationRequest(
             replicate_configuration=replicate_configuration
         )
+
+    @staticmethod
+    def convert_function_to_function_schema(f: Function) -> schema_types.FunctionSchema:
+        function_schema = schema_types.FunctionSchema(
+            name=f.name,
+            description=f.description,
+            type=f.type,
+            input_field_names=f.input_field_names,
+            output_field_names=f.output_field_names,
+        )
+        for k, v in f.params.items():
+            kv_pair = common_types.KeyValuePair(key=str(k), value=str(v))
+            function_schema.params.append(kv_pair)
+        return function_schema

pymilvus/grpc_gen/common_pb2.py

@@ -241,6 +241,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__ = ()
@@ -354,9 +355,7 @@ class ObjectPrivilege(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     PrivilegeAddFileResource: _ClassVar[ObjectPrivilege]
     PrivilegeRemoveFileResource: _ClassVar[ObjectPrivilege]
     PrivilegeListFileResources: _ClassVar[ObjectPrivilege]
-    PrivilegeAddCollectionFunction: _ClassVar[ObjectPrivilege]
-    PrivilegeAlterCollectionFunction: _ClassVar[ObjectPrivilege]
-    PrivilegeDropCollectionFunction: _ClassVar[ObjectPrivilege]
+    PrivilegeUpdateReplicateConfiguration: _ClassVar[ObjectPrivilege]
 
 class StateCode(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     __slots__ = ()
@@ -608,6 +607,7 @@ ListDatabases: MsgType
 AlterDatabase: MsgType
 DescribeDatabase: MsgType
 AddCollectionField: MsgType
+AlterWAL: MsgType
 Dsl: DslType
 BoolExprV1: DslType
 UndefiedState: CompactionState
@@ -703,9 +703,7 @@ PrivilegeAddCollectionField: ObjectPrivilege
 PrivilegeAddFileResource: ObjectPrivilege
 PrivilegeRemoveFileResource: ObjectPrivilege
 PrivilegeListFileResources: ObjectPrivilege
-PrivilegeAddCollectionFunction: ObjectPrivilege
-PrivilegeAlterCollectionFunction: ObjectPrivilege
-PrivilegeDropCollectionFunction: ObjectPrivilege
+PrivilegeUpdateReplicateConfiguration: ObjectPrivilege
 Initializing: StateCode
 Healthy: StateCode
 Abnormal: StateCode

pymilvus/grpc_gen/common_pb2.pyi

@@ -158,7 +158,7 @@ class ListAliasesResponse(_message.Message):
     def __init__(self, status: _Optional[_Union[_common_pb2.Status, _Mapping]] = ..., db_name: _Optional[str] = ..., collection_name: _Optional[str] = ..., aliases: _Optional[_Iterable[str]] = ...) -> None: ...
 
 class CreateCollectionRequest(_message.Message):
-    __slots__ = ("base", "db_name", "collection_name", "schema", "shards_num", "consistency_level", "properties", "num_partitions")
+    __slots__ = ("base", "db_name", "collection_name", "schema", "shards_num", "consistency_level", "properties", "num_partitions", "external_source", "external_spec")
     BASE_FIELD_NUMBER: _ClassVar[int]
     DB_NAME_FIELD_NUMBER: _ClassVar[int]
     COLLECTION_NAME_FIELD_NUMBER: _ClassVar[int]
@@ -167,6 +167,8 @@ class CreateCollectionRequest(_message.Message):
     CONSISTENCY_LEVEL_FIELD_NUMBER: _ClassVar[int]
     PROPERTIES_FIELD_NUMBER: _ClassVar[int]
     NUM_PARTITIONS_FIELD_NUMBER: _ClassVar[int]
+    EXTERNAL_SOURCE_FIELD_NUMBER: _ClassVar[int]
+    EXTERNAL_SPEC_FIELD_NUMBER: _ClassVar[int]
     base: _common_pb2.MsgBase
     db_name: str
     collection_name: str
@@ -175,7 +177,9 @@ class CreateCollectionRequest(_message.Message):
     consistency_level: _common_pb2.ConsistencyLevel
     properties: _containers.RepeatedCompositeFieldContainer[_common_pb2.KeyValuePair]
     num_partitions: int
-    def __init__(self, base: _Optional[_Union[_common_pb2.MsgBase, _Mapping]] = ..., db_name: _Optional[str] = ..., collection_name: _Optional[str] = ..., schema: _Optional[bytes] = ..., shards_num: _Optional[int] = ..., consistency_level: _Optional[_Union[_common_pb2.ConsistencyLevel, str]] = ..., properties: _Optional[_Iterable[_Union[_common_pb2.KeyValuePair, _Mapping]]] = ..., num_partitions: _Optional[int] = ...) -> None: ...
+    external_source: str
+    external_spec: str
+    def __init__(self, base: _Optional[_Union[_common_pb2.MsgBase, _Mapping]] = ..., db_name: _Optional[str] = ..., collection_name: _Optional[str] = ..., schema: _Optional[bytes] = ..., shards_num: _Optional[int] = ..., consistency_level: _Optional[_Union[_common_pb2.ConsistencyLevel, str]] = ..., properties: _Optional[_Iterable[_Union[_common_pb2.KeyValuePair, _Mapping]]] = ..., num_partitions: _Optional[int] = ..., external_source: _Optional[str] = ..., external_spec: _Optional[str] = ...) -> None: ...
 
 class DropCollectionRequest(_message.Message):
     __slots__ = ("base", "db_name", "collection_name")
@@ -1436,7 +1440,7 @@ class LoadBalanceRequest(_message.Message):
     def __init__(self, base: _Optional[_Union[_common_pb2.MsgBase, _Mapping]] = ..., src_nodeID: _Optional[int] = ..., dst_nodeIDs: _Optional[_Iterable[int]] = ..., sealed_segmentIDs: _Optional[_Iterable[int]] = ..., collectionName: _Optional[str] = ..., db_name: _Optional[str] = ...) -> None: ...
 
 class ManualCompactionRequest(_message.Message):
-    __slots__ = ("collectionID", "timetravel", "majorCompaction", "collection_name", "db_name", "partition_id", "channel", "segment_ids", "l0Compaction")
+    __slots__ = ("collectionID", "timetravel", "majorCompaction", "collection_name", "db_name", "partition_id", "channel", "segment_ids", "l0Compaction", "target_size")
     COLLECTIONID_FIELD_NUMBER: _ClassVar[int]
     TIMETRAVEL_FIELD_NUMBER: _ClassVar[int]
     MAJORCOMPACTION_FIELD_NUMBER: _ClassVar[int]
@@ -1446,6 +1450,7 @@ class ManualCompactionRequest(_message.Message):
     CHANNEL_FIELD_NUMBER: _ClassVar[int]
     SEGMENT_IDS_FIELD_NUMBER: _ClassVar[int]
     L0COMPACTION_FIELD_NUMBER: _ClassVar[int]
+    TARGET_SIZE_FIELD_NUMBER: _ClassVar[int]
     collectionID: int
     timetravel: int
     majorCompaction: bool
@@ -1455,7 +1460,8 @@ class ManualCompactionRequest(_message.Message):
     channel: str
     segment_ids: _containers.RepeatedScalarFieldContainer[int]
     l0Compaction: bool
-    def __init__(self, collectionID: _Optional[int] = ..., timetravel: _Optional[int] = ..., majorCompaction: bool = ..., collection_name: _Optional[str] = ..., db_name: _Optional[str] = ..., partition_id: _Optional[int] = ..., channel: _Optional[str] = ..., segment_ids: _Optional[_Iterable[int]] = ..., l0Compaction: bool = ...) -> None: ...
+    target_size: int
+    def __init__(self, collectionID: _Optional[int] = ..., timetravel: _Optional[int] = ..., majorCompaction: bool = ..., collection_name: _Optional[str] = ..., db_name: _Optional[str] = ..., partition_id: _Optional[int] = ..., channel: _Optional[str] = ..., segment_ids: _Optional[_Iterable[int]] = ..., l0Compaction: bool = ..., target_size: _Optional[int] = ...) -> None: ...
 
 class ManualCompactionResponse(_message.Message):
     __slots__ = ("status", "compactionID", "compactionPlanCount")