refactor(api): migrate dataset endpoints to BaseModel (#37957)

Co-authored-by: Asuka Minato <i@asukaminato.eu.org>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
chariri
2026-07-09 04:22:30 +00:00
committed by GitHub
co-authored by Asuka Minato autofix-ci[bot]
parent 9bb3b1fa98
commit 89d5f74a40
22 changed files with 1871 additions and 1403 deletions
+23 -23
View File
@@ -30,6 +30,7 @@ from controllers.console.wraps import (
with_current_tenant_id,
with_current_user,
)
from core.entities.knowledge_entities import IndexingEstimate
from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
from core.indexing_runner import IndexingRunner
from core.plugin.impl.model_runtime_factory import create_plugin_provider_manager
@@ -268,21 +269,10 @@ class ErrorDocsResponse(DocumentStatusListResponse):
total: int
class IndexingEstimatePreviewItemResponse(ResponseModel):
content: str
child_chunks: list[str] | None = None
summary: str | None = None
class IndexingEstimateQaPreviewItemResponse(ResponseModel):
question: str
answer: str
class IndexingEstimateResponse(ResponseModel):
total_segments: int
preview: list[IndexingEstimatePreviewItemResponse]
qa_preview: list[IndexingEstimateQaPreviewItemResponse] | None = None
class IndexingEstimateResponse(IndexingEstimate):
tokens: int
total_price: float | int
currency: str
class RetrievalSettingResponse(ResponseModel):
@@ -647,7 +637,7 @@ class DatasetApi(Resource):
else:
data["embedding_available"] = True
return data, 200
return dump_response(DatasetDetailWithPartialMembersResponse, data), 200
@console_ns.doc("update_dataset")
@console_ns.doc(description="Update dataset details")
@@ -717,7 +707,7 @@ class DatasetApi(Resource):
partial_member_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session())
result_data.update({"partial_member_list": partial_member_list})
return result_data, 200
return dump_response(DatasetDetailWithPartialMembersResponse, result_data), 200
@setup_required
@login_required
@@ -760,7 +750,7 @@ class DatasetUseCheckApi(Resource):
dataset_id_str = str(dataset_id)
dataset_is_using = DatasetService.dataset_use_check(dataset_id_str, db.session())
return {"is_using": dataset_is_using}, 200
return UsageCheckResponse(is_using=dataset_is_using).model_dump(mode="json"), 200
@console_ns.route("/datasets/<uuid:dataset_id>/queries")
@@ -901,7 +891,17 @@ class DatasetIndexingEstimateApi(Resource):
except Exception as e:
raise IndexingEstimateError(str(e))
return response.model_dump(), 200
return (
IndexingEstimateResponse(
tokens=0,
total_price=0,
currency="USD",
total_segments=response.total_segments,
preview=response.preview,
qa_preview=response.qa_preview,
).model_dump(mode="json", exclude_none=True),
200,
)
@console_ns.route("/datasets/<uuid:dataset_id>/related-apps")
@@ -1018,7 +1018,7 @@ class DatasetApiKeyApi(Resource):
keys = db.session.scalars(
select(ApiToken).where(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_tenant_id)
).all()
return ApiKeyList.model_validate({"data": keys}, from_attributes=True).model_dump(mode="json")
return dump_response(ApiKeyList, {"data": keys})
@console_ns.response(200, "API key created successfully", console_ns.models[ApiKeyItem.__name__])
@console_ns.response(400, "Maximum keys exceeded")
@@ -1052,7 +1052,7 @@ class DatasetApiKeyApi(Resource):
api_token.type = self.resource_type
db.session.add(api_token)
db.session.commit()
return ApiKeyItem.model_validate(api_token, from_attributes=True).model_dump(mode="json"), 200
return dump_response(ApiKeyItem, api_token), 200
@console_ns.route("/datasets/api-keys/<uuid:api_key_id>")
@@ -1107,7 +1107,7 @@ class DatasetEnableApiApi(Resource):
DatasetService.update_dataset_api_status(dataset_id_str, status == "enable", db.session())
return {"result": "success"}, 200
return SimpleResultResponse(result="success").model_dump(mode="json"), 200
@console_ns.route("/datasets/api-base-info")
@@ -1120,7 +1120,7 @@ class DatasetApiBaseUrlApi(Resource):
@account_initialization_required
def get(self):
base = dify_config.SERVICE_API_URL or request.host_url.rstrip("/")
return {"api_base_url": normalize_api_base_url(base)}
return ApiBaseUrlResponse(api_base_url=normalize_api_base_url(base)).model_dump(mode="json")
@console_ns.route("/datasets/retrieval-setting")
@@ -10,16 +10,17 @@ from uuid import UUID
import sqlalchemy as sa
from flask import request, send_file
from flask_restx import Resource
from pydantic import BaseModel, Field, RootModel, field_validator
from pydantic import BaseModel, Field, JsonValue, field_validator
from sqlalchemy import asc, desc, func, select
from werkzeug.exceptions import Forbidden, NotFound
import services
from controllers.common.controller_schemas import DocumentBatchDownloadZipPayload
from controllers.common.fields import BinaryFileResponse, SimpleResultMessageResponse, SimpleResultResponse, UrlResponse
from controllers.common.fields import SimpleResultMessageResponse, SimpleResultResponse, UrlResponse
from controllers.common.schema import register_response_schema_models, register_schema_models
from controllers.console import console_ns
from controllers.console.wraps import RBACPermission, RBACResourceScope, rbac_permission_required
from core.entities.knowledge_entities import IndexingEstimate
from core.errors.error import (
LLMBadRequestError,
ModelCurrentlyNotSupportError,
@@ -29,6 +30,7 @@ from core.errors.error import (
from core.indexing_runner import IndexingRunner
from core.model_manager import ModelManager
from core.plugin.impl.exc import PluginDaemonClientSideError
from core.rag.entities import Rule
from core.rag.extractor.entity.datasource_type import DatasourceType
from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo, WebsiteInfo
from core.rag.index_processor.constant.index_type import IndexTechniqueType
@@ -49,7 +51,7 @@ from libs.login import login_required
from libs.pagination import paginate_query
from models import Account, DatasetProcessRule, Document, DocumentSegment, UploadFile
from models.dataset import DocumentPipelineExecutionLog
from models.enums import IndexingStatus, SegmentStatus
from models.enums import IndexingStatus, ProcessRuleMode, SegmentStatus
from services.dataset_ref_service import DatasetRefService
from services.dataset_service import DatasetService, DocumentService
from services.entities.knowledge_entities.knowledge_entities import KnowledgeConfig, ProcessRule, RetrievalModel
@@ -148,8 +150,91 @@ class DocumentWithSegmentsListResponse(ResponseModel):
page: int
class OpaqueObjectResponse(RootModel[dict[str, Any]]):
root: dict[str, Any]
class IndexingEstimateResponse(IndexingEstimate):
tokens: int
total_price: float | int
currency: str
class DocumentDetailResponse(ResponseModel):
id: str
position: int | None = None
data_source_type: str | None = None
data_source_info: Any = None
data_source_detail_dict: Any = None
dataset_process_rule_id: str | None = None
dataset_process_rule: Any = None
document_process_rule: Any = None
name: str | None = None
created_from: str | None = None
created_by: str | None = None
created_at: int | None = None
tokens: int | None = None
indexing_status: str | None = None
completed_at: int | None = None
updated_at: int | None = None
indexing_latency: float | None = None
error: str | None = None
enabled: bool | None = None
disabled_at: int | None = None
disabled_by: str | None = None
archived: bool | None = None
doc_type: str | None = None
doc_metadata: list[DocumentMetadataResponse] | None = None
segment_count: int | None = None
average_segment_length: float | None = None
hit_count: int | None = None
display_status: str | None = None
doc_form: str | None = None
doc_language: str | None = None
need_summary: bool | None = None
@field_validator("data_source_type", "indexing_status", "display_status", "doc_form", mode="before")
@classmethod
def _normalize_enum_fields(cls, value: Any) -> Any:
return normalize_enum(value)
class SummaryStatusResponse(ResponseModel):
completed: int = 0
generating: int = 0
error: int = 0
not_started: int = 0
timeout: int = 0
class SummaryEntryResponse(ResponseModel):
segment_id: str
segment_position: int
status: str
summary_preview: str | None = None
error: str | None = None
created_at: int | None = None
updated_at: int | None = None
@field_validator("status", mode="before")
@classmethod
def _normalize_status(cls, value: Any) -> Any:
return normalize_enum(value)
class DocumentSummaryStatusResponse(ResponseModel):
total_segments: int
summary_status: SummaryStatusResponse
summaries: list[SummaryEntryResponse]
class ProcessRuleResponse(ResponseModel):
mode: ProcessRuleMode
rules: Rule | None = None
limits: dict[str, Any]
class DocumentPipelineExecutionLogResponse(ResponseModel):
datasource_info: JsonValue | None = None
datasource_type: str | None = None
input_data: JsonValue | None = None
datasource_node_id: str | None = None
register_schema_models(
@@ -165,7 +250,6 @@ register_schema_models(
)
register_response_schema_models(
console_ns,
BinaryFileResponse,
SimpleResultMessageResponse,
SimpleResultResponse,
UrlResponse,
@@ -175,7 +259,11 @@ register_response_schema_models(
DocumentWithSegmentsResponse,
DatasetAndDocumentResponse,
DocumentWithSegmentsListResponse,
OpaqueObjectResponse,
IndexingEstimateResponse,
DocumentDetailResponse,
DocumentSummaryStatusResponse,
ProcessRuleResponse,
DocumentPipelineExecutionLogResponse,
)
@@ -225,7 +313,7 @@ class GetProcessRuleApi(Resource):
@console_ns.doc("get_process_rule")
@console_ns.doc(description="Get dataset document processing rules")
@console_ns.doc(params={"document_id": "Document ID (optional)"})
@console_ns.response(200, "Process rules retrieved successfully", console_ns.models[OpaqueObjectResponse.__name__])
@console_ns.response(200, "Process rules retrieved successfully", console_ns.models[ProcessRuleResponse.__name__])
@setup_required
@login_required
@account_initialization_required
@@ -264,7 +352,7 @@ class GetProcessRuleApi(Resource):
mode = dataset_process_rule.mode
rules = dataset_process_rule.rules_dict
return {"mode": mode, "rules": rules, "limits": limits}
return dump_response(ProcessRuleResponse, {"mode": mode, "rules": rules, "limits": limits})
@console_ns.route("/datasets/<uuid:dataset_id>/documents")
@@ -491,7 +579,7 @@ class DatasetInitApi(Resource):
@console_ns.doc(description="Initialize dataset with documents")
@console_ns.expect(console_ns.models[KnowledgeConfig.__name__])
@console_ns.response(
201, "Dataset initialized successfully", console_ns.models[DatasetAndDocumentResponse.__name__]
200, "Dataset initialized successfully", console_ns.models[DatasetAndDocumentResponse.__name__]
)
@console_ns.response(400, "Invalid request parameters")
@setup_required
@@ -557,7 +645,7 @@ class DocumentIndexingEstimateApi(DocumentResource):
@console_ns.response(
200,
"Indexing estimate calculated successfully",
console_ns.models[OpaqueObjectResponse.__name__],
console_ns.models[IndexingEstimateResponse.__name__],
)
@console_ns.response(404, "Document not found")
@console_ns.response(400, "Document already finished")
@@ -578,8 +666,6 @@ class DocumentIndexingEstimateApi(DocumentResource):
data_process_rule = document.dataset_process_rule
data_process_rule_dict = data_process_rule.to_dict() if data_process_rule else {}
response = {"tokens": 0, "total_price": 0, "currency": "USD", "total_segments": 0, "preview": []}
if document.data_source_type == "upload_file":
data_source_info = document.data_source_info_dict
if data_source_info and "upload_file_id" in data_source_info:
@@ -610,7 +696,18 @@ class DocumentIndexingEstimateApi(DocumentResource):
"English",
dataset_id_str,
)
return estimate_response.model_dump(), 200
return (
# TODO: why using zero here? the same for the below endpoint
IndexingEstimateResponse(
tokens=0,
total_price=0,
currency="USD",
total_segments=estimate_response.total_segments,
preview=estimate_response.preview,
qa_preview=estimate_response.qa_preview,
).model_dump(mode="json", exclude_none=True),
200,
)
except LLMBadRequestError:
raise ProviderNotInitializeError(
"No Embedding Model available. Please configure a valid provider "
@@ -623,15 +720,24 @@ class DocumentIndexingEstimateApi(DocumentResource):
except Exception as e:
raise IndexingEstimateError(str(e))
return response, 200
return (
IndexingEstimateResponse(
tokens=0,
total_price=0,
currency="USD",
total_segments=0,
preview=[],
).model_dump(mode="json", exclude_none=True),
200,
)
@console_ns.route("/datasets/<uuid:dataset_id>/batch/<string:batch>/indexing-estimate")
class DocumentBatchIndexingEstimateApi(DocumentResource):
@console_ns.response(
200,
"Batch indexing estimate calculated successfully",
console_ns.models[OpaqueObjectResponse.__name__],
"Indexing estimate calculated successfully",
console_ns.models[IndexingEstimateResponse.__name__],
)
@setup_required
@login_required
@@ -643,7 +749,16 @@ class DocumentBatchIndexingEstimateApi(DocumentResource):
dataset_id_str = str(dataset_id)
documents = self.get_batch_documents(dataset_id_str, batch, current_user)
if not documents:
return {"tokens": 0, "total_price": 0, "currency": "USD", "total_segments": 0, "preview": []}, 200
return (
IndexingEstimateResponse(
tokens=0,
total_price=0,
currency="USD",
total_segments=0,
preview=[],
).model_dump(mode="json", exclude_none=True),
200,
)
data_process_rule = documents[0].dataset_process_rule
data_process_rule_dict = data_process_rule.to_dict() if data_process_rule else {}
extract_settings = []
@@ -717,7 +832,17 @@ class DocumentBatchIndexingEstimateApi(DocumentResource):
"English",
dataset_id_str,
)
return response.model_dump(), 200
return (
IndexingEstimateResponse(
tokens=0,
total_price=0,
currency="USD",
total_segments=response.total_segments,
preview=response.preview,
qa_preview=response.qa_preview,
).model_dump(mode="json", exclude_none=True),
200,
)
except LLMBadRequestError:
raise ProviderNotInitializeError(
"No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider."
@@ -854,7 +979,7 @@ class DocumentApi(DocumentResource):
"metadata": "Metadata inclusion (all/only/without)",
}
)
@console_ns.response(200, "Document retrieved successfully", console_ns.models[OpaqueObjectResponse.__name__])
@console_ns.response(200, "Document retrieved successfully", console_ns.models[DocumentDetailResponse.__name__])
@console_ns.response(404, "Document not found")
@setup_required
@login_required
@@ -871,46 +996,21 @@ class DocumentApi(DocumentResource):
if metadata not in self.METADATA_CHOICES:
raise InvalidMetadataError(f"Invalid metadata value: {metadata}")
metadata_fields = {"doc_type", "doc_metadata"}
if metadata == "only":
response = {"id": document.id, "doc_type": document.doc_type, "doc_metadata": document.doc_metadata_details}
elif metadata == "without":
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, db.session())
document_process_rules = document.dataset_process_rule.to_dict() if document.dataset_process_rule else {}
response = {
"id": document.id,
"position": document.position,
"data_source_type": document.data_source_type,
"data_source_info": document.data_source_info_dict,
"data_source_detail_dict": document.data_source_detail_dict,
"dataset_process_rule_id": document.dataset_process_rule_id,
"dataset_process_rule": dataset_process_rules,
"document_process_rule": document_process_rules,
"name": document.name,
"created_from": document.created_from,
"created_by": document.created_by,
"created_at": int(document.created_at.timestamp()),
"tokens": document.tokens,
"indexing_status": document.indexing_status,
"completed_at": int(document.completed_at.timestamp()) if document.completed_at else None,
"updated_at": int(document.updated_at.timestamp()) if document.updated_at else None,
"indexing_latency": document.indexing_latency,
"error": document.error,
"enabled": document.enabled,
"disabled_at": int(document.disabled_at.timestamp()) if document.disabled_at else None,
"disabled_by": document.disabled_by,
"archived": document.archived,
"segment_count": document.segment_count,
"average_segment_length": document.average_segment_length,
"hit_count": document.hit_count,
"display_status": document.display_status,
"doc_form": document.doc_form,
"doc_language": document.doc_language,
"need_summary": document.need_summary if document.need_summary is not None else False,
}
else:
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, db.session())
document_process_rules = document.dataset_process_rule.to_dict() if document.dataset_process_rule else {}
response = {
response = DocumentDetailResponse.model_validate(
{
"id": document.id,
"doc_type": document.doc_type,
"doc_metadata": document.doc_metadata_details,
}
)
return response.model_dump(mode="json", include={"id", *metadata_fields}, exclude_unset=True), 200
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, db.session())
document_process_rules = document.dataset_process_rule.to_dict() if document.dataset_process_rule else {}
response = DocumentDetailResponse.model_validate(
{
"id": document.id,
"position": document.position,
"data_source_type": document.data_source_type,
@@ -943,8 +1043,9 @@ class DocumentApi(DocumentResource):
"doc_language": document.doc_language,
"need_summary": document.need_summary if document.need_summary is not None else False,
}
return response, 200
)
exclude = metadata_fields if metadata == "without" else None
return response.model_dump(mode="json", exclude=exclude, exclude_unset=True), 200
@setup_required
@login_required
@@ -990,7 +1091,9 @@ class DocumentDownloadApi(DocumentResource):
def get(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID) -> dict[str, Any]:
# Reuse the shared permission/tenant checks implemented in DocumentResource.
document = self.get_document(str(dataset_id), str(document_id), current_user, current_tenant_id)
return {"url": DocumentService.get_document_download_url(document, db.session())}
return UrlResponse(url=DocumentService.get_document_download_url(document, db.session())).model_dump(
mode="json"
)
@console_ns.route("/datasets/<uuid:dataset_id>/documents/download-zip")
@@ -999,7 +1102,7 @@ class DocumentBatchDownloadZipApi(DocumentResource):
@console_ns.doc("download_dataset_documents_as_zip")
@console_ns.doc(description="Download selected dataset documents as a single ZIP archive (upload-file only)")
@console_ns.response(200, "ZIP archive generated successfully", console_ns.models[BinaryFileResponse.__name__])
@console_ns.response(200, "ZIP archive downloaded successfully")
@setup_required
@login_required
@account_initialization_required
@@ -1034,6 +1137,7 @@ class DocumentBatchDownloadZipApi(DocumentResource):
)
cleanup = stack.pop_all()
response.call_on_close(cleanup.close)
# response-contract:ignore binary ZIP download response
return response
@@ -1093,7 +1197,7 @@ class DocumentProcessingApi(DocumentResource):
document.is_paused = False
db.session.commit()
return {"result": "success"}, 200
return SimpleResultResponse(result="success").model_dump(mode="json"), 200
@console_ns.route("/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/metadata")
@@ -1152,7 +1256,9 @@ class DocumentMetadataApi(DocumentResource):
document.updated_at = naive_utc_now()
db.session.commit()
return {"result": "success", "message": "Document metadata updated."}, 200
return SimpleResultMessageResponse(result="success", message="Document metadata updated.").model_dump(
mode="json"
), 200
@console_ns.route("/datasets/<uuid:dataset_id>/documents/status/<string:action>/batch")
@@ -1194,7 +1300,7 @@ class DocumentStatusApi(DocumentResource):
except NotFound as e:
raise NotFound(str(e))
return {"result": "success"}, 200
return SimpleResultResponse(result="success").model_dump(mode="json"), 200
@console_ns.route("/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/processing/pause")
@@ -1321,7 +1427,7 @@ class DocumentRenameApi(DocumentResource):
# The role of the current user in the ta table must be admin, owner, editor, or dataset_operator
if not current_user.is_dataset_editor:
raise Forbidden()
dataset = DatasetService.get_dataset(dataset_id, db.session())
dataset = DatasetService.get_dataset(str(dataset_id), db.session())
if not dataset:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_operator_permission(current_user, dataset, session=db.session())
@@ -1363,15 +1469,15 @@ class WebsiteDocumentSyncApi(DocumentResource):
# sync document
DocumentService.sync_website_document(dataset_id_str, document, db.session())
return {"result": "success"}, 200
return SimpleResultResponse(result="success").model_dump(mode="json"), 200
@console_ns.route("/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/pipeline-execution-log")
class DocumentPipelineExecutionLogApi(DocumentResource):
@console_ns.response(
200,
"Document pipeline execution log retrieved successfully",
console_ns.models[OpaqueObjectResponse.__name__],
"Pipeline execution log retrieved successfully",
console_ns.models[DocumentPipelineExecutionLogResponse.__name__],
)
@setup_required
@login_required
@@ -1394,18 +1500,16 @@ class DocumentPipelineExecutionLogApi(DocumentResource):
.limit(1)
)
if not log:
return {
"datasource_info": None,
"datasource_type": None,
"input_data": None,
"datasource_node_id": None,
}, 200
return {
"datasource_info": json.loads(log.datasource_info),
"datasource_type": log.datasource_type,
"input_data": log.input_data,
"datasource_node_id": log.datasource_node_id,
}, 200
return DocumentPipelineExecutionLogResponse().model_dump(mode="json"), 200
return dump_response(
DocumentPipelineExecutionLogResponse,
{
"datasource_info": json.loads(log.datasource_info),
"datasource_type": log.datasource_type,
"input_data": log.input_data,
"datasource_node_id": log.datasource_node_id,
},
), 200
@console_ns.route("/datasets/<uuid:dataset_id>/documents/generate-summary")
@@ -1508,7 +1612,7 @@ class DocumentGenerateSummaryApi(Resource):
dataset_id_str,
)
return {"result": "success"}, 200
return SimpleResultResponse(result="success").model_dump(mode="json"), 200
@console_ns.route("/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/summary-status")
@@ -1516,7 +1620,11 @@ class DocumentSummaryStatusApi(DocumentResource):
@console_ns.doc("get_document_summary_status")
@console_ns.doc(description="Get summary index generation status for a document")
@console_ns.doc(params={"dataset_id": "Dataset ID", "document_id": "Document ID"})
@console_ns.response(200, "Summary status retrieved successfully", console_ns.models[OpaqueObjectResponse.__name__])
@console_ns.response(
200,
"Summary status retrieved successfully",
console_ns.models[DocumentSummaryStatusResponse.__name__],
)
@console_ns.response(404, "Document not found")
@setup_required
@login_required
@@ -1534,6 +1642,7 @@ class DocumentSummaryStatusApi(DocumentResource):
- generating: Number of summaries being generated
- error: Number of summaries with errors
- not_started: Number of segments without summary records
- timeout: Number of summaries that timed out
- summaries: List of summary records with status and content preview
"""
dataset_id_str = str(dataset_id)
@@ -1559,4 +1668,4 @@ class DocumentSummaryStatusApi(DocumentResource):
session=db.session(),
)
return result, 200
return dump_response(DocumentSummaryStatusResponse, result), 200
@@ -391,7 +391,7 @@ class DatasetDocumentSegmentApi(Resource):
SegmentService.update_segments_status(segment_ids, action, dataset, document, db.session())
except Exception as e:
raise InvalidActionError(str(e))
return dump_response(SimpleResultResponse, {"result": "success"}), 200
return SimpleResultResponse(result="success").model_dump(mode="json"), 200
@console_ns.route("/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segment")
+89 -107
View File
@@ -1,20 +1,16 @@
from datetime import datetime
from typing import Any
from uuid import UUID
from flask import request
from flask_restx import Resource, fields, marshal
from pydantic import BaseModel, Field, RootModel
from flask_restx import Resource
from pydantic import AliasChoices, BaseModel, Field, field_validator
from sqlalchemy.orm import Session
from werkzeug.exceptions import Forbidden, InternalServerError, NotFound
import services
from controllers.common.fields import UsageCountResponse
from controllers.common.schema import (
get_or_create_model,
query_params_from_model,
register_response_schema_models,
register_schema_models,
)
from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
from controllers.console import console_ns
from controllers.console.app.wraps import with_session
from controllers.console.datasets.error import DatasetNameDuplicateError
@@ -28,21 +24,9 @@ from controllers.console.wraps import (
with_current_tenant_id,
with_current_user,
)
from extensions.ext_database import db
from fields.base import ResponseModel
from fields.dataset_fields import (
dataset_detail_fields,
dataset_retrieval_model_fields,
doc_metadata_fields,
external_knowledge_info_fields,
external_retrieval_model_fields,
icon_info_fields,
keyword_setting_fields,
reranking_model_fields,
tag_fields,
vector_setting_fields,
weighted_score_fields,
)
from fields.dataset_fields import DatasetDetailResponse
from libs.helper import dump_response
from libs.login import login_required
from models import Account
from services.dataset_service import DatasetService
@@ -51,50 +35,10 @@ from services.external_knowledge_service import ExternalDatasetService
from services.hit_testing_service import HitTestingService
from services.knowledge_service import BedrockRetrievalSetting, ExternalDatasetTestService
register_response_schema_models(console_ns, UsageCountResponse)
def _build_dataset_detail_model():
keyword_setting_model = get_or_create_model("DatasetKeywordSetting", keyword_setting_fields)
vector_setting_model = get_or_create_model("DatasetVectorSetting", vector_setting_fields)
weighted_score_fields_copy = weighted_score_fields.copy()
weighted_score_fields_copy["keyword_setting"] = fields.Nested(keyword_setting_model)
weighted_score_fields_copy["vector_setting"] = fields.Nested(vector_setting_model)
weighted_score_model = get_or_create_model("DatasetWeightedScore", weighted_score_fields_copy)
reranking_model = get_or_create_model("DatasetRerankingModel", reranking_model_fields)
dataset_retrieval_model_fields_copy = dataset_retrieval_model_fields.copy()
dataset_retrieval_model_fields_copy["reranking_model"] = fields.Nested(reranking_model)
dataset_retrieval_model_fields_copy["weights"] = fields.Nested(weighted_score_model, allow_null=True)
dataset_retrieval_model = get_or_create_model("DatasetRetrievalModel", dataset_retrieval_model_fields_copy)
tag_model = get_or_create_model("Tag", tag_fields)
doc_metadata_model = get_or_create_model("DatasetDocMetadata", doc_metadata_fields)
external_knowledge_info_model = get_or_create_model("ExternalKnowledgeInfo", external_knowledge_info_fields)
external_retrieval_model = get_or_create_model("ExternalRetrievalModel", external_retrieval_model_fields)
icon_info_model = get_or_create_model("DatasetIconInfo", icon_info_fields)
dataset_detail_fields_copy = dataset_detail_fields.copy()
dataset_detail_fields_copy["retrieval_model_dict"] = fields.Nested(dataset_retrieval_model)
dataset_detail_fields_copy["tags"] = fields.List(fields.Nested(tag_model))
dataset_detail_fields_copy["external_knowledge_info"] = fields.Nested(external_knowledge_info_model)
dataset_detail_fields_copy["external_retrieval_model"] = fields.Nested(external_retrieval_model, allow_null=True)
dataset_detail_fields_copy["doc_metadata"] = fields.List(fields.Nested(doc_metadata_model))
dataset_detail_fields_copy["icon_info"] = fields.Nested(icon_info_model)
return get_or_create_model("DatasetDetail", dataset_detail_fields_copy)
try:
dataset_detail_model = console_ns.models["DatasetDetail"]
except KeyError:
dataset_detail_model = _build_dataset_detail_model()
class ExternalKnowledgeApiPayload(BaseModel):
name: str = Field(..., min_length=1, max_length=40)
settings: dict[str, object]
settings: dict[str, Any]
class ExternalDatasetCreatePayload(BaseModel):
@@ -102,15 +46,13 @@ class ExternalDatasetCreatePayload(BaseModel):
external_knowledge_id: str
name: str = Field(..., min_length=1, max_length=100)
description: str | None = Field(None, max_length=400)
external_retrieval_model: dict[str, object] | None = Field(default=None)
external_retrieval_model: dict[str, Any] | None = None
class ExternalHitTestingPayload(BaseModel):
query: str
external_retrieval_model: dict[str, object] | None = Field(default=None)
metadata_filtering_conditions: dict[str, object] | None = Field(
default=None,
)
external_retrieval_model: dict[str, Any] | None = None
metadata_filtering_conditions: dict[str, Any] | None = None
class BedrockRetrievalPayload(BaseModel):
@@ -125,7 +67,7 @@ class ExternalApiTemplateListQuery(BaseModel):
keyword: str | None = Field(default=None, description="Search keyword")
class ExternalKnowledgeDatasetBindingResponse(ResponseModel):
class ExternalKnowledgeApiBindingResponse(ResponseModel):
id: str
name: str
@@ -135,22 +77,52 @@ class ExternalKnowledgeApiResponse(ResponseModel):
tenant_id: str
name: str
description: str
settings: dict[str, Any] | None = Field(default=None)
dataset_bindings: list[ExternalKnowledgeDatasetBindingResponse] = Field(default_factory=list)
settings: dict[str, Any] | None = Field(validation_alias=AliasChoices("settings_dict", "settings"))
dataset_bindings: list[ExternalKnowledgeApiBindingResponse]
created_by: str
created_at: str
@field_validator("created_at", mode="before")
@classmethod
def _normalize_created_at(cls, value: datetime | str) -> str:
if isinstance(value, datetime):
return value.isoformat()
return value
class ExternalKnowledgeApiListResponse(ResponseModel):
data: list[ExternalKnowledgeApiResponse]
has_more: bool
limit: int
total: int
total: int | None
page: int
class ExternalRetrievalTestResponse(RootModel[dict[str, Any] | list[dict[str, Any]]]):
root: dict[str, Any] | list[dict[str, Any]]
class ExternalHitTestingQueryResponse(ResponseModel):
content: str
class ExternalHitTestingRecordResponse(ResponseModel):
content: str | None = None
title: str | None = None
score: float | None = None
metadata: dict[str, Any] | None = None
class ExternalHitTestingResponse(ResponseModel):
query: ExternalHitTestingQueryResponse
records: list[ExternalHitTestingRecordResponse]
class BedrockRetrievalRecordResponse(ResponseModel):
metadata: dict[str, Any] | None = None
score: float
title: str | None = None
content: str | None = None
class BedrockRetrievalResponse(ResponseModel):
records: list[BedrockRetrievalRecordResponse]
register_schema_models(
@@ -163,9 +135,16 @@ register_schema_models(
)
register_response_schema_models(
console_ns,
UsageCountResponse,
DatasetDetailResponse,
ExternalKnowledgeApiBindingResponse,
ExternalKnowledgeApiResponse,
ExternalKnowledgeApiListResponse,
ExternalRetrievalTestResponse,
ExternalHitTestingQueryResponse,
ExternalHitTestingRecordResponse,
ExternalHitTestingResponse,
BedrockRetrievalRecordResponse,
BedrockRetrievalResponse,
)
@@ -189,24 +168,26 @@ class ExternalApiTemplateListApi(Resource):
external_knowledge_apis, total = ExternalDatasetService.get_external_knowledge_apis(
query.page, query.limit, current_tenant_id, query.keyword
)
response = {
"data": [item.to_dict() for item in external_knowledge_apis],
"has_more": len(external_knowledge_apis) == query.limit,
"limit": query.limit,
"total": total,
"page": query.page,
}
return response, 200
return ExternalKnowledgeApiListResponse(
data=[ExternalKnowledgeApiResponse.model_validate(item) for item in external_knowledge_apis],
has_more=len(external_knowledge_apis) == query.limit,
limit=query.limit,
total=total,
page=query.page,
).model_dump(mode="json"), 200
@setup_required
@login_required
@account_initialization_required
@console_ns.doc("create_external_api_template")
@console_ns.doc(description="Create external knowledge API template")
@console_ns.expect(console_ns.models[ExternalKnowledgeApiPayload.__name__])
@console_ns.response(
201,
"External API template created successfully",
console_ns.models[ExternalKnowledgeApiResponse.__name__],
)
@console_ns.response(403, "Permission denied")
@setup_required
@login_required
@account_initialization_required
@with_current_user
@with_current_tenant_id
@with_session
@@ -229,7 +210,7 @@ class ExternalApiTemplateListApi(Resource):
except services.errors.dataset.DatasetNameDuplicateError:
raise DatasetNameDuplicateError()
return external_knowledge_api.to_dict(), 201
return dump_response(ExternalKnowledgeApiResponse, external_knowledge_api), 201
@console_ns.route("/datasets/external-knowledge-api/<uuid:external_knowledge_api_id>")
@@ -256,17 +237,21 @@ class ExternalApiTemplateApi(Resource):
if external_knowledge_api is None:
raise NotFound("API template not found.")
return external_knowledge_api.to_dict(), 200
return dump_response(ExternalKnowledgeApiResponse, external_knowledge_api), 200
@console_ns.doc("update_external_api_template")
@console_ns.doc(description="Update external knowledge API template")
@console_ns.doc(params={"external_knowledge_api_id": "External knowledge API ID"})
@console_ns.expect(console_ns.models[ExternalKnowledgeApiPayload.__name__])
@console_ns.response(
200,
"External API template updated successfully",
console_ns.models[ExternalKnowledgeApiResponse.__name__],
)
@console_ns.response(404, "Template not found")
@setup_required
@login_required
@account_initialization_required
@console_ns.expect(console_ns.models[ExternalKnowledgeApiPayload.__name__])
@with_current_user
@with_current_tenant_id
@with_session
@@ -284,7 +269,7 @@ class ExternalApiTemplateApi(Resource):
session=session,
)
return external_knowledge_api.to_dict(), 200
return dump_response(ExternalKnowledgeApiResponse, external_knowledge_api), 200
@setup_required
@login_required
@@ -322,7 +307,7 @@ class ExternalApiUseCheckApi(Resource):
external_knowledge_api_is_using, count = ExternalDatasetService.external_knowledge_api_use_check(
external_knowledge_api_id_str, current_tenant_id, session=session
)
return {"is_using": external_knowledge_api_is_using, "count": count}, 200
return UsageCountResponse(is_using=external_knowledge_api_is_using, count=count).model_dump(mode="json"), 200
@console_ns.route("/datasets/external")
@@ -330,7 +315,9 @@ class ExternalDatasetCreateApi(Resource):
@console_ns.doc("create_external_dataset")
@console_ns.doc(description="Create external knowledge dataset")
@console_ns.expect(console_ns.models[ExternalDatasetCreatePayload.__name__])
@console_ns.response(201, "External dataset created successfully", dataset_detail_model)
@console_ns.response(
201, "External dataset created successfully", console_ns.models[DatasetDetailResponse.__name__]
)
@console_ns.response(400, "Invalid parameters")
@console_ns.response(403, "Permission denied")
@setup_required
@@ -360,17 +347,16 @@ class ExternalDatasetCreateApi(Resource):
except services.errors.dataset.DatasetNameDuplicateError:
raise DatasetNameDuplicateError()
item = marshal(dataset, dataset_detail_fields)
dataset_id_str = item["id"]
dataset_id_str = str(dataset.id)
permission_keys_map = enterprise_rbac_service.RBACService.DatasetPermissions.batch_get(
str(current_tenant_id),
current_user.id,
[dataset_id_str],
session=session,
)
item["permission_keys"] = permission_keys_map.get(dataset_id_str, [])
return item, 201
data = DatasetDetailResponse.model_validate(dataset).model_dump(mode="json")
data["permission_keys"] = permission_keys_map.get(dataset_id_str, [])
return data, 201
@console_ns.route("/datasets/<uuid:dataset_id>/external-hit-testing")
@@ -382,7 +368,7 @@ class ExternalKnowledgeHitTestingApi(Resource):
@console_ns.response(
200,
"External hit testing completed successfully",
console_ns.models[ExternalRetrievalTestResponse.__name__],
console_ns.models[ExternalHitTestingResponse.__name__],
)
@console_ns.response(404, "Dataset not found")
@console_ns.response(400, "Invalid parameters")
@@ -394,12 +380,12 @@ class ExternalKnowledgeHitTestingApi(Resource):
@with_session
def post(self, session: Session, current_user: Account, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
@@ -416,7 +402,7 @@ class ExternalKnowledgeHitTestingApi(Resource):
metadata_filtering_conditions=payload.metadata_filtering_conditions,
)
return response
return dump_response(ExternalHitTestingResponse, response)
except Exception as e:
raise InternalServerError(str(e))
@@ -427,11 +413,7 @@ class BedrockRetrievalApi(Resource):
@console_ns.doc("bedrock_retrieval_test")
@console_ns.doc(description="Bedrock retrieval test (internal use only)")
@console_ns.expect(console_ns.models[BedrockRetrievalPayload.__name__])
@console_ns.response(
200,
"Bedrock retrieval test completed",
console_ns.models[ExternalRetrievalTestResponse.__name__],
)
@console_ns.response(200, "Bedrock retrieval test completed", console_ns.models[BedrockRetrievalResponse.__name__])
def post(self):
payload = BedrockRetrievalPayload.model_validate(console_ns.payload or {})
@@ -439,4 +421,4 @@ class BedrockRetrievalApi(Resource):
result = ExternalDatasetTestService.knowledge_retrieval(
payload.retrieval_setting, payload.query, payload.knowledge_id
)
return result, 200
return dump_response(BedrockRetrievalResponse, result), 200