add import linter

api/.ruff.toml

@@ -49,6 +49,7 @@ select = [
     "S301", # suspicious-pickle-usage, disallow use of `pickle` and its wrappers.
     "S302", # suspicious-marshal-usage, disallow use of `marshal` module
     "S311", # suspicious-non-cryptographic-random-usage,
+    "TID",   # flake8-tidy-imports
 
 ]
 
@@ -84,6 +85,7 @@ ignore = [
     "SIM113",  # enumerate-for-loop
     "SIM117",  # multiple-with-statements
     "SIM210",  # if-expr-with-true-false
+    "TID252",  # allow relative imports from parent modules
 ]
 
 [lint.per-file-ignores]
@@ -112,3 +114,11 @@ allowed-unused-imports = [
     "tests.integration_tests",
     "tests.unit_tests",
 ]
+
+[lint.flake8-tidy-imports]
+
+[lint.flake8-tidy-imports.banned-api."flask_restx.reqparse"]
+msg = "Use Pydantic payload/query models instead of reqparse."
+
+[lint.flake8-tidy-imports.banned-api."flask_restx.reqparse.RequestParser"]
+msg = "Use Pydantic payload/query models instead of reqparse."

api/controllers/console/datasets/hit_testing_base.py

@@ -1,7 +1,7 @@
 import logging
 from typing import Any
 
-from flask_restx import marshal, reqparse
+from flask_restx import marshal
 from pydantic import BaseModel, Field
 from werkzeug.exceptions import Forbidden, InternalServerError, NotFound
 
@@ -56,15 +56,10 @@ class DatasetsHitTestingBase:
         HitTestingService.hit_testing_args_check(args)
 
     @staticmethod
-    def parse_args():
-        parser = (
-            reqparse.RequestParser()
-            .add_argument("query", type=str, required=False, location="json")
-            .add_argument("attachment_ids", type=list, required=False, location="json")
-            .add_argument("retrieval_model", type=dict, required=False, location="json")
-            .add_argument("external_retrieval_model", type=dict, required=False, location="json")
-        )
-        return parser.parse_args()
+    def parse_args(payload: dict[str, Any] | None = None) -> dict[str, Any]:
+        """Validate and return hit-testing arguments from an incoming payload."""
+        hit_testing_payload = HitTestingPayload.model_validate(payload or {})
+        return hit_testing_payload.model_dump(exclude_none=True)
 
     @staticmethod
     def perform_hit_testing(dataset, args):

api/controllers/service_api/dataset/hit_testing.py

@@ -1,7 +1,10 @@
-from controllers.console.datasets.hit_testing_base import DatasetsHitTestingBase
+from controllers.common.schema import register_schema_model
+from controllers.console.datasets.hit_testing_base import DatasetsHitTestingBase, HitTestingPayload
 from controllers.service_api import service_api_ns
 from controllers.service_api.wraps import DatasetApiResource, cloud_edition_billing_rate_limit_check
 
+register_schema_model(service_api_ns, HitTestingPayload)
+
 
 @service_api_ns.route("/datasets/<uuid:dataset_id>/hit-testing", "/datasets/<uuid:dataset_id>/retrieve")
 class HitTestingApi(DatasetApiResource, DatasetsHitTestingBase):
@@ -15,6 +18,7 @@ class HitTestingApi(DatasetApiResource, DatasetsHitTestingBase):
             404: "Dataset not found",
         }
     )
+    @service_api_ns.expect(service_api_ns.models[HitTestingPayload.__name__])
     @cloud_edition_billing_rate_limit_check("knowledge", "dataset")
     def post(self, tenant_id, dataset_id):
         """Perform hit testing on a dataset.
@@ -24,7 +28,7 @@ class HitTestingApi(DatasetApiResource, DatasetsHitTestingBase):
         dataset_id_str = str(dataset_id)
 
         dataset = self.get_and_validate_dataset(dataset_id_str)
-        args = self.parse_args()
+        args = self.parse_args(service_api_ns.payload)
         self.hit_testing_args_check(args)
 
         return self.perform_hit_testing(dataset, args)