refactor(api): migrate console/service_api.dataset.hit_testing to BaseModel (#36533)

api/controllers/console/datasets/hit_testing.py

@@ -1,15 +1,12 @@
 from __future__ import annotations
 
-from datetime import datetime
-from typing import Any
 from uuid import UUID
 
 from flask_restx import Resource
-from pydantic import Field, field_validator
 
-from controllers.common.schema import register_schema_models
-from fields.base import ResponseModel
-from libs.helper import to_timestamp
+from controllers.common.schema import register_response_schema_models, register_schema_models
+from fields.hit_testing_fields import HitTestingResponse
+from libs.helper import dump_response
 from libs.login import login_required
 
 from .. import console_ns
@@ -20,86 +17,8 @@ from ..wraps import (
     setup_required,
 )
 
-
-class HitTestingDocument(ResponseModel):
-    id: str | None = None
-    data_source_type: str | None = None
-    name: str | None = None
-    doc_type: str | None = None
-    doc_metadata: Any | None = None
-
-
-class HitTestingSegment(ResponseModel):
-    id: str | None = None
-    position: int | None = None
-    document_id: str | None = None
-    content: str | None = None
-    sign_content: str | None = None
-    answer: str | None = None
-    word_count: int | None = None
-    tokens: int | None = None
-    keywords: list[str] = Field(default_factory=list)
-    index_node_id: str | None = None
-    index_node_hash: str | None = None
-    hit_count: int | None = None
-    enabled: bool | None = None
-    disabled_at: int | None = None
-    disabled_by: str | None = None
-    status: str | None = None
-    created_by: str | None = None
-    created_at: int | None = None
-    indexing_at: int | None = None
-    completed_at: int | None = None
-    error: str | None = None
-    stopped_at: int | None = None
-    document: HitTestingDocument | None = None
-
-    @field_validator("disabled_at", "created_at", "indexing_at", "completed_at", "stopped_at", mode="before")
-    @classmethod
-    def _normalize_timestamp(cls, value: datetime | int | None) -> int | None:
-        return to_timestamp(value)
-
-
-class HitTestingChildChunk(ResponseModel):
-    id: str | None = None
-    content: str | None = None
-    position: int | None = None
-    score: float | None = None
-
-
-class HitTestingFile(ResponseModel):
-    id: str | None = None
-    name: str | None = None
-    size: int | None = None
-    extension: str | None = None
-    mime_type: str | None = None
-    source_url: str | None = None
-
-
-class HitTestingRecord(ResponseModel):
-    segment: HitTestingSegment | None = None
-    child_chunks: list[HitTestingChildChunk] = Field(default_factory=list)
-    score: float | None = None
-    tsne_position: Any | None = None
-    files: list[HitTestingFile] = Field(default_factory=list)
-    summary: str | None = None
-
-
-class HitTestingResponse(ResponseModel):
-    query: str
-    records: list[HitTestingRecord] = Field(default_factory=list)
-
-
-register_schema_models(
-    console_ns,
-    HitTestingPayload,
-    HitTestingDocument,
-    HitTestingSegment,
-    HitTestingChildChunk,
-    HitTestingFile,
-    HitTestingRecord,
-    HitTestingResponse,
-)
+register_schema_models(console_ns, HitTestingPayload)
+register_response_schema_models(console_ns, HitTestingResponse)
 
 
 @console_ns.route("/datasets/<uuid:dataset_id>/hit-testing")
@@ -119,12 +38,11 @@ class HitTestingApi(Resource, DatasetsHitTestingBase):
     @login_required
     @account_initialization_required
     @cloud_edition_billing_rate_limit_check("knowledge")
-    def post(self, dataset_id: UUID):
+    def post(self, dataset_id: UUID) -> dict[str, object]:
         dataset_id_str = str(dataset_id)
 
         dataset = self.get_and_validate_dataset(dataset_id_str)
-        payload = HitTestingPayload.model_validate(console_ns.payload or {})
-        args = payload.model_dump(exclude_none=True)
+        args = self.parse_args(console_ns.payload)
         self.hit_testing_args_check(args)
 
-        return HitTestingResponse.model_validate(self.perform_hit_testing(dataset, args)).model_dump(mode="json")
+        return dump_response(HitTestingResponse, self.perform_hit_testing(dataset, args))

api/controllers/console/datasets/hit_testing_base.py

@@ -1,7 +1,6 @@
 import logging
-from typing import Any
+from typing import Any, cast
 
-from flask_restx import marshal
 from pydantic import BaseModel, Field
 from werkzeug.exceptions import Forbidden, InternalServerError, NotFound
 
@@ -19,10 +18,10 @@ from core.errors.error import (
     ProviderTokenNotInitError,
     QuotaExceededError,
 )
-from fields.hit_testing_fields import hit_testing_record_fields
 from graphon.model_runtime.errors.invoke import InvokeError
 from libs.login import current_user
 from models.account import Account
+from models.dataset import Dataset
 from services.dataset_service import DatasetService
 from services.entities.knowledge_entities.knowledge_entities import RetrievalModel
 from services.hit_testing_service import HitTestingService
@@ -38,16 +37,6 @@ class HitTestingPayload(BaseModel):
 
 
 class DatasetsHitTestingBase:
-    @staticmethod
-    def _extract_hit_testing_query(query: Any) -> str:
-        """Return the query string from the service response shape."""
-        if isinstance(query, dict):
-            content = query.get("content")
-            if isinstance(content, str):
-                return content
-
-        raise ValueError("Invalid hit testing query response")
-
     @staticmethod
     def _prepare_hit_testing_records(records: Any) -> list[dict[str, Any]]:
         """Ensure collection fields match the API schema before response validation."""
@@ -63,6 +52,7 @@ class DatasetsHitTestingBase:
             segment = normalized_record.get("segment")
             if isinstance(segment, dict):
                 normalized_segment = dict(segment)
+                normalized_segment.setdefault("sign_content", None)
                 if normalized_segment.get("keywords") is None:
                     normalized_segment["keywords"] = []
                 normalized_record["segment"] = normalized_segment
@@ -73,12 +63,15 @@ class DatasetsHitTestingBase:
             if normalized_record.get("files") is None:
                 normalized_record["files"] = []
 
+            normalized_record.setdefault("tsne_position", None)
+            normalized_record.setdefault("summary", None)
+
             normalized_records.append(normalized_record)
 
         return normalized_records
 
     @staticmethod
-    def get_and_validate_dataset(dataset_id: str):
+    def get_and_validate_dataset(dataset_id: str) -> Dataset:
         assert isinstance(current_user, Account)
         dataset = DatasetService.get_dataset(dataset_id)
         if dataset is None:
@@ -92,33 +85,35 @@ class DatasetsHitTestingBase:
         return dataset
 
     @staticmethod
-    def hit_testing_args_check(args: dict[str, Any]):
+    def hit_testing_args_check(args: dict[str, Any]) -> None:
         HitTestingService.hit_testing_args_check(args)
 
     @staticmethod
-    def parse_args(payload: dict[str, Any]) -> dict[str, Any]:
+    def parse_args(payload: dict[str, Any] | 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):
+    def perform_hit_testing(dataset: Dataset, args: dict[str, Any]) -> dict[str, Any]:
         assert isinstance(current_user, Account)
         try:
             response = HitTestingService.retrieve(
                 dataset=dataset,
-                query=args.get("query"),
+                query=cast(str, args.get("query")),
                 account=current_user,
                 retrieval_model=args.get("retrieval_model"),
-                external_retrieval_model=args.get("external_retrieval_model"),
+                external_retrieval_model=cast(dict[str, Any], args.get("external_retrieval_model")),
                 attachment_ids=args.get("attachment_ids"),
                 limit=10,
             )
+            query = response.get("query")
+            if not isinstance(query, dict) or not isinstance(query.get("content"), str):
+                raise ValueError("Invalid hit testing query response")
+
             return {
-                "query": DatasetsHitTestingBase._extract_hit_testing_query(response.get("query")),
-                "records": DatasetsHitTestingBase._prepare_hit_testing_records(
-                    marshal(response.get("records", []), hit_testing_record_fields)
-                ),
+                "query": {"content": query["content"]},
+                "records": DatasetsHitTestingBase._prepare_hit_testing_records(response.get("records", [])),
             }
         except services.errors.index.IndexNotInitializedError:
             raise DatasetNotInitializedError()

api/controllers/service_api/dataset/hit_testing.py

@@ -1,11 +1,14 @@
 from uuid import UUID
 
-from controllers.common.schema import register_schema_model
+from controllers.common.schema import register_response_schema_models, register_schema_models
 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
+from fields.hit_testing_fields import HitTestingResponse
+from libs.helper import dump_response
 
-register_schema_model(service_api_ns, HitTestingPayload)
+register_schema_models(service_api_ns, HitTestingPayload)
+register_response_schema_models(service_api_ns, HitTestingResponse)
 
 
 @service_api_ns.route("/datasets/<uuid:dataset_id>/hit-testing", "/datasets/<uuid:dataset_id>/retrieve")
@@ -13,16 +16,16 @@ class HitTestingApi(DatasetApiResource, DatasetsHitTestingBase):
     @service_api_ns.doc("dataset_hit_testing")
     @service_api_ns.doc(description="Perform hit testing on a dataset")
     @service_api_ns.doc(params={"dataset_id": "Dataset ID"})
-    @service_api_ns.doc(
-        responses={
-            200: "Hit testing results",
-            401: "Unauthorized - invalid API token",
-            404: "Dataset not found",
-        }
+    @service_api_ns.response(
+        200,
+        "Hit testing results",
+        model=service_api_ns.models[HitTestingResponse.__name__],
     )
+    @service_api_ns.response(401, "Unauthorized - invalid API token")
+    @service_api_ns.response(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: UUID):
+    def post(self, tenant_id: str, dataset_id: UUID) -> dict[str, object]:
         """Perform hit testing on a dataset.
 
         Tests retrieval performance for the specified dataset.
@@ -33,4 +36,4 @@ class HitTestingApi(DatasetApiResource, DatasetsHitTestingBase):
         args = self.parse_args(service_api_ns.payload)
         self.hit_testing_args_check(args)
 
-        return self.perform_hit_testing(dataset, args)
+        return dump_response(HitTestingResponse, self.perform_hit_testing(dataset, args))