mirror of
https://github.com/langgenius/dify.git
synced 2026-09-24 23:22:26 +08:00
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:
co-authored by
Asuka Minato
autofix-ci[bot]
parent
9bb3b1fa98
commit
89d5f74a40
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user