refactor(api): Finish the pydantic.BaseModel migration

api/controllers/common/fields.py

@@ -1,14 +1,13 @@
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, Literal
 
 from pydantic import BaseModel, ConfigDict, Field, RootModel, computed_field
 
 from fields.base import ResponseModel
 from graphon.file import helpers as file_helpers
 from models.model import IconType
 
-type JSONValue = str | int | float | bool | None | dict[str, Any] | list[Any]
 type JSONObject = dict[str, Any]
 
 
@@ -24,10 +23,6 @@ class SimpleResultResponse(ResponseModel):
     result: str
 
 
-class GeneratedAppResponse(RootModel[JSONValue]):
-    root: JSONValue
-
-
 class EventStreamResponse(RootModel[str]):
     root: str
 
@@ -52,6 +47,11 @@ class AudioTranscriptResponse(ResponseModel):
     text: str
 
 
+class ValidationResultResponse(ResponseModel):
+    result: Literal["success", "error"]
+    error: str | None = None
+
+
 class SimpleResultMessageResponse(ResponseModel):
     result: str
     message: str

api/controllers/console/app/agent_app_feature.py

@@ -30,7 +30,6 @@
 )
 from events.app_event import app_model_config_was_updated
 from extensions.ext_database import db
-from libs.helper import dump_response
 from libs.login import login_required
 from models import Account
 from models.agent_config_entities import (
@@ -99,4 +98,4 @@ def post(self, tenant_id: str, current_user: Account, agent_id: UUID):
 
         app_model_config_was_updated.send(app_model, app_model_config=new_app_model_config)
 
-        return dump_response(SimpleResultResponse, {"result": "success"})
+        return SimpleResultResponse(result="success").model_dump(mode="json")

api/controllers/console/app/annotation.py

@@ -1,7 +1,7 @@
 from typing import Any, Literal
 from uuid import UUID
 
-from flask import abort, make_response, request
+from flask import abort, request
 from flask_restx import Resource
 from pydantic import BaseModel, Field, TypeAdapter, field_validator
 
@@ -26,10 +26,12 @@
     AnnotationExportList,
     AnnotationHitHistory,
     AnnotationHitHistoryList,
+    AnnotationJobStatusDetailResponse,
+    AnnotationJobStatusResponse,
     AnnotationList,
 )
 from fields.base import ResponseModel
-from libs.helper import uuid_value
+from libs.helper import dump_response, uuid_value
 from libs.login import login_required
 from services.annotation_service import (
     AppAnnotationService,
@@ -99,23 +101,23 @@ def validate_message_id(cls, value: str) -> str:
         return uuid_value(value)
 
 
-class AnnotationJobStatusResponse(ResponseModel):
-    job_id: str | None = None
-    job_status: str | None = None
-    error_msg: str | None = None
-    record_count: int | None = None
-
-
-class AnnotationEmbeddingModelResponse(ResponseModel):
+class AnnotationSettingEmbeddingModelResponse(ResponseModel):
     embedding_provider_name: str | None = None
     embedding_model_name: str | None = None
 
 
 class AnnotationSettingResponse(ResponseModel):
-    id: str | None = None
     enabled: bool
+    id: str | None = None
     score_threshold: float | None = None
-    embedding_model: AnnotationEmbeddingModelResponse | None = None
+    embedding_model: AnnotationSettingEmbeddingModelResponse | None = None
+
+
+class AnnotationBatchImportResponse(ResponseModel):
+    job_id: str | None = None
+    job_status: str | None = None
+    record_count: int | None = None
+    error_msg: str | None = None
 
 
 register_schema_models(
@@ -142,7 +144,10 @@ class AnnotationSettingResponse(ResponseModel):
     AnnotationHitHistory,
     AnnotationHitHistoryList,
     AnnotationJobStatusResponse,
+    AnnotationJobStatusDetailResponse,
+    AnnotationSettingEmbeddingModelResponse,
     AnnotationSettingResponse,
+    AnnotationBatchImportResponse,
 )
 
 
@@ -172,7 +177,7 @@ def post(self, app_id: UUID, action: Literal["enable", "disable"]):
                 result = AppAnnotationService.enable_app_annotation(enable_args, str(app_id))
             case "disable":
                 result = AppAnnotationService.disable_app_annotation(str(app_id))
-        return result, 200
+        return dump_response(AnnotationJobStatusResponse, result), 200
 
 
 @console_ns.route("/apps/<uuid:app_id>/annotation-setting")
@@ -193,7 +198,7 @@ class AppAnnotationSettingDetailApi(Resource):
     @rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
     def get(self, app_id: UUID):
         result = AppAnnotationService.get_app_annotation_setting_by_app_id(str(app_id))
-        return result, 200
+        return dump_response(AnnotationSettingResponse, result), 200
 
 
 @console_ns.route("/apps/<uuid:app_id>/annotation-settings/<uuid:annotation_setting_id>")
@@ -218,7 +223,7 @@ def post(self, app_id: UUID, annotation_setting_id: UUID):
         result = AppAnnotationService.update_app_annotation_setting(
             str(app_id), annotation_setting_id_str, setting_args
         )
-        return result, 200
+        return dump_response(AnnotationSettingResponse, result), 200
 
 
 @console_ns.route("/apps/<uuid:app_id>/annotation-reply/<string:action>/status/<uuid:job_id>")
@@ -227,9 +232,7 @@ class AnnotationReplyActionStatusApi(Resource):
     @console_ns.doc(description="Get status of annotation reply action job")
     @console_ns.doc(params={"app_id": "Application ID", "job_id": "Job ID", "action": "Action type"})
     @console_ns.response(
-        200,
-        "Job status retrieved successfully",
-        console_ns.models[AnnotationJobStatusResponse.__name__],
+        200, "Job status retrieved successfully", console_ns.models[AnnotationJobStatusDetailResponse.__name__]
     )
     @console_ns.response(403, "Insufficient permissions")
     @setup_required
@@ -251,7 +254,9 @@ def get(self, app_id: UUID, job_id: UUID, action: str):
             app_annotation_error_key = f"{action}_app_annotation_error_{job_id_str}"
             error_msg = redis_client.get(app_annotation_error_key).decode()
 
-        return {"job_id": job_id_str, "job_status": job_status, "error_msg": error_msg}, 200
+        return AnnotationJobStatusDetailResponse(
+            job_id=job_id_str, job_status=job_status, error_msg=error_msg
+        ).model_dump(mode="json"), 200
 
 
 @console_ns.route("/apps/<uuid:app_id>/annotations")
@@ -275,14 +280,9 @@ def get(self, app_id: UUID):
 
         annotation_list, total = AppAnnotationService.get_annotation_list_by_app_id(str(app_id), page, limit, keyword)
         annotation_models = TypeAdapter(list[Annotation]).validate_python(annotation_list, from_attributes=True)
-        response = AnnotationList(
-            data=annotation_models,
-            has_more=len(annotation_list) == limit,
-            limit=limit,
-            total=total,
-            page=page,
-        )
-        return response.model_dump(mode="json"), 200
+        return AnnotationList(
+            data=annotation_models, has_more=len(annotation_list) == limit, limit=limit, total=total, page=page
+        ).model_dump(mode="json"), 200
 
     @console_ns.doc("create_annotation")
     @console_ns.doc(description="Create a new annotation for an app")
@@ -308,7 +308,7 @@ def post(self, app_id: UUID):
         if args.question is not None:
             upsert_args["question"] = args.question
         annotation = AppAnnotationService.up_insert_app_annotation_from_message(upsert_args, str(app_id))
-        return Annotation.model_validate(annotation, from_attributes=True).model_dump(mode="json")
+        return dump_response(Annotation, annotation), 201
 
     @setup_required
     @login_required
@@ -357,14 +357,14 @@ class AnnotationExportApi(Resource):
     def get(self, app_id: UUID):
         annotation_list = AppAnnotationService.export_annotation_list_by_app_id(str(app_id))
         annotation_models = TypeAdapter(list[Annotation]).validate_python(annotation_list, from_attributes=True)
-        response_data = AnnotationExportList(data=annotation_models).model_dump(mode="json")
-
-        # Create response with secure headers for CSV export
-        response = make_response(response_data, 200)
-        response.headers["Content-Type"] = "application/json; charset=utf-8"
-        response.headers["X-Content-Type-Options"] = "nosniff"
-
-        return response
+        return (
+            AnnotationExportList(data=annotation_models).model_dump(mode="json"),
+            200,
+            {
+                "Content-Type": "application/json; charset=utf-8",
+                "X-Content-Type-Options": "nosniff",
+            },
+        )
 
 
 @console_ns.route("/apps/<uuid:app_id>/annotations/<uuid:annotation_id>")
@@ -392,7 +392,7 @@ def post(self, app_id: UUID, annotation_id: UUID):
         annotation = AppAnnotationService.update_app_annotation_directly(
             update_args, str(app_id), str(annotation_id), db.session
         )
-        return Annotation.model_validate(annotation, from_attributes=True).model_dump(mode="json")
+        return dump_response(Annotation, annotation)
 
     @setup_required
     @login_required
@@ -411,9 +411,7 @@ class AnnotationBatchImportApi(Resource):
     @console_ns.doc(description="Batch import annotations from CSV file with rate limiting and security checks")
     @console_ns.doc(params={"app_id": "Application ID"})
     @console_ns.response(
-        200,
-        "Batch import started successfully",
-        console_ns.models[AnnotationJobStatusResponse.__name__],
+        200, "Batch import started successfully", console_ns.models[AnnotationBatchImportResponse.__name__]
     )
     @console_ns.response(403, "Insufficient permissions")
     @console_ns.response(400, "No file uploaded or too many files")
@@ -460,7 +458,10 @@ def post(self, app_id: UUID):
         if file_size == 0:
             raise ValueError("The uploaded file is empty")
 
-        return AppAnnotationService.batch_import_app_annotations(str(app_id), file)
+        return dump_response(
+            AnnotationBatchImportResponse,
+            AppAnnotationService.batch_import_app_annotations(str(app_id), file),
+        )
 
 
 @console_ns.route("/apps/<uuid:app_id>/annotations/batch-import-status/<uuid:job_id>")
@@ -469,9 +470,7 @@ class AnnotationBatchImportStatusApi(Resource):
     @console_ns.doc(description="Get status of batch import job")
     @console_ns.doc(params={"app_id": "Application ID", "job_id": "Job ID"})
     @console_ns.response(
-        200,
-        "Job status retrieved successfully",
-        console_ns.models[AnnotationJobStatusResponse.__name__],
+        200, "Job status retrieved successfully", console_ns.models[AnnotationJobStatusDetailResponse.__name__]
     )
     @console_ns.response(403, "Insufficient permissions")
     @setup_required
@@ -491,7 +490,9 @@ def get(self, app_id: UUID, job_id: UUID):
             indexing_error_msg_key = f"app_annotation_batch_import_error_msg_{str(job_id)}"
             error_msg = redis_client.get(indexing_error_msg_key).decode()
 
-        return {"job_id": job_id, "job_status": job_status, "error_msg": error_msg}, 200
+        return AnnotationJobStatusDetailResponse(
+            job_id=str(job_id), job_status=job_status, error_msg=error_msg
+        ).model_dump(mode="json"), 200
 
 
 @console_ns.route("/apps/<uuid:app_id>/annotations/<uuid:annotation_id>/hit-histories")
@@ -520,11 +521,6 @@ def get(self, app_id: UUID, annotation_id: UUID):
         history_models = TypeAdapter(list[AnnotationHitHistory]).validate_python(
             annotation_hit_history_list, from_attributes=True
         )
-        response = AnnotationHitHistoryList(
-            data=history_models,
-            has_more=len(annotation_hit_history_list) == limit,
-            limit=limit,
-            total=total,
-            page=page,
-        )
-        return response.model_dump(mode="json")
+        return AnnotationHitHistoryList(
+            data=history_models, has_more=len(annotation_hit_history_list) == limit, limit=limit, total=total, page=page
+        ).model_dump(mode="json")