add typing to all wraps (#25405)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/core/rag/datasource/vdb/matrixone/matrixone_vector.py

@@ -1,8 +1,9 @@
 import json
 import logging
 import uuid
+from collections.abc import Callable
 from functools import wraps
-from typing import Any, Optional
+from typing import Any, Concatenate, Optional, ParamSpec, TypeVar
 
 from mo_vector.client import MoVectorClient  # type: ignore
 from pydantic import BaseModel, model_validator
@@ -17,7 +18,6 @@ from extensions.ext_redis import redis_client
 from models.dataset import Dataset
 
 logger = logging.getLogger(__name__)
-from typing import ParamSpec, TypeVar
 
 P = ParamSpec("P")
 R = TypeVar("R")
@@ -47,16 +47,6 @@ class MatrixoneConfig(BaseModel):
         return values
 
 
-def ensure_client(func):
-    @wraps(func)
-    def wrapper(self, *args, **kwargs):
-        if self.client is None:
-            self.client = self._get_client(None, False)
-        return func(self, *args, **kwargs)
-
-    return wrapper
-
-
 class MatrixoneVector(BaseVector):
     """
     Matrixone vector storage implementation.
@@ -216,6 +206,19 @@ class MatrixoneVector(BaseVector):
         self.client.delete()
 
 
+T = TypeVar("T", bound=MatrixoneVector)
+
+
+def ensure_client(func: Callable[Concatenate[T, P], R]):
+    @wraps(func)
+    def wrapper(self: T, *args: P.args, **kwargs: P.kwargs):
+        if self.client is None:
+            self.client = self._get_client(None, False)
+        return func(self, *args, **kwargs)
+
+    return wrapper
+
+
 class MatrixoneVectorFactory(AbstractVectorFactory):
     def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> MatrixoneVector:
         if dataset.index_struct_dict: