Merge remote-tracking branch 'upstream/main' into feat/rag-2

api/controllers/console/datasets/datasets_document.py

@@ -1,16 +1,17 @@
 import json
 import logging
 from argparse import ArgumentTypeError
+from collections.abc import Sequence
 from typing import Literal, cast
 
 from flask import request
 from flask_login import current_user
-from flask_restx import Resource, marshal, marshal_with, reqparse
+from flask_restx import Resource, fields, marshal, marshal_with, reqparse
 from sqlalchemy import asc, desc, select
 from werkzeug.exceptions import Forbidden, NotFound
 
 import services
-from controllers.console import api
+from controllers.console import api, console_ns
 from controllers.console.app.error import (
     ProviderModelCurrentlyNotSupportError,
     ProviderNotInitializeError,
@@ -81,7 +82,7 @@ class DocumentResource(Resource):
 
         return document
 
-    def get_batch_documents(self, dataset_id: str, batch: str) -> list[Document]:
+    def get_batch_documents(self, dataset_id: str, batch: str) -> Sequence[Document]:
         dataset = DatasetService.get_dataset(dataset_id)
         if not dataset:
             raise NotFound("Dataset not found.")
@@ -99,7 +100,12 @@ class DocumentResource(Resource):
         return documents
 
 
+@console_ns.route("/datasets/process-rule")
 class GetProcessRuleApi(Resource):
+    @api.doc("get_process_rule")
+    @api.doc(description="Get dataset document processing rules")
+    @api.doc(params={"document_id": "Document ID (optional)"})
+    @api.response(200, "Process rules retrieved successfully")
     @setup_required
     @login_required
     @account_initialization_required
@@ -141,7 +147,21 @@ class GetProcessRuleApi(Resource):
         return {"mode": mode, "rules": rules, "limits": limits}
 
 
+@console_ns.route("/datasets/<uuid:dataset_id>/documents")
 class DatasetDocumentListApi(Resource):
+    @api.doc("get_dataset_documents")
+    @api.doc(description="Get documents in a dataset")
+    @api.doc(
+        params={
+            "dataset_id": "Dataset ID",
+            "page": "Page number (default: 1)",
+            "limit": "Number of items per page (default: 20)",
+            "keyword": "Search keyword",
+            "sort": "Sort order (default: -created_at)",
+            "fetch": "Fetch full details (default: false)",
+        }
+    )
+    @api.response(200, "Documents retrieved successfully")
     @setup_required
     @login_required
     @account_initialization_required
@@ -325,7 +345,23 @@ class DatasetDocumentListApi(Resource):
         return {"result": "success"}, 204
 
 
+@console_ns.route("/datasets/init")
 class DatasetInitApi(Resource):
+    @api.doc("init_dataset")
+    @api.doc(description="Initialize dataset with documents")
+    @api.expect(
+        api.model(
+            "DatasetInitRequest",
+            {
+                "upload_file_id": fields.String(required=True, description="Upload file ID"),
+                "indexing_technique": fields.String(description="Indexing technique"),
+                "process_rule": fields.Raw(description="Processing rules"),
+                "data_source": fields.Raw(description="Data source configuration"),
+            },
+        )
+    )
+    @api.response(201, "Dataset initialized successfully", dataset_and_document_fields)
+    @api.response(400, "Invalid request parameters")
     @setup_required
     @login_required
     @account_initialization_required
@@ -395,7 +431,14 @@ class DatasetInitApi(Resource):
         return response
 
 
+@console_ns.route("/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/indexing-estimate")
 class DocumentIndexingEstimateApi(DocumentResource):
+    @api.doc("estimate_document_indexing")
+    @api.doc(description="Estimate document indexing cost")
+    @api.doc(params={"dataset_id": "Dataset ID", "document_id": "Document ID"})
+    @api.response(200, "Indexing estimate calculated successfully")
+    @api.response(404, "Document not found")
+    @api.response(400, "Document already finished")
     @setup_required
     @login_required
     @account_initialization_required
@@ -595,7 +638,13 @@ class DocumentBatchIndexingStatusApi(DocumentResource):
         return data
 
 
+@console_ns.route("/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/indexing-status")
 class DocumentIndexingStatusApi(DocumentResource):
+    @api.doc("get_document_indexing_status")
+    @api.doc(description="Get document indexing status")
+    @api.doc(params={"dataset_id": "Dataset ID", "document_id": "Document ID"})
+    @api.response(200, "Indexing status retrieved successfully")
+    @api.response(404, "Document not found")
     @setup_required
     @login_required
     @account_initialization_required
@@ -637,9 +686,21 @@ class DocumentIndexingStatusApi(DocumentResource):
         return marshal(document_dict, document_status_fields)
 
 
+@console_ns.route("/datasets/<uuid:dataset_id>/documents/<uuid:document_id>")
 class DocumentApi(DocumentResource):
     METADATA_CHOICES = {"all", "only", "without"}
 
+    @api.doc("get_document")
+    @api.doc(description="Get document details")
+    @api.doc(
+        params={
+            "dataset_id": "Dataset ID",
+            "document_id": "Document ID",
+            "metadata": "Metadata inclusion (all/only/without)",
+        }
+    )
+    @api.response(200, "Document retrieved successfully")
+    @api.response(404, "Document not found")
     @setup_required
     @login_required
     @account_initialization_required
@@ -748,7 +809,16 @@ class DocumentApi(DocumentResource):
         return {"result": "success"}, 204
 
 
+@console_ns.route("/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/processing/<string:action>")
 class DocumentProcessingApi(DocumentResource):
+    @api.doc("update_document_processing")
+    @api.doc(description="Update document processing status (pause/resume)")
+    @api.doc(
+        params={"dataset_id": "Dataset ID", "document_id": "Document ID", "action": "Action to perform (pause/resume)"}
+    )
+    @api.response(200, "Processing status updated successfully")
+    @api.response(404, "Document not found")
+    @api.response(400, "Invalid action")
     @setup_required
     @login_required
     @account_initialization_required
@@ -783,7 +853,23 @@ class DocumentProcessingApi(DocumentResource):
         return {"result": "success"}, 200
 
 
+@console_ns.route("/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/metadata")
 class DocumentMetadataApi(DocumentResource):
+    @api.doc("update_document_metadata")
+    @api.doc(description="Update document metadata")
+    @api.doc(params={"dataset_id": "Dataset ID", "document_id": "Document ID"})
+    @api.expect(
+        api.model(
+            "UpdateDocumentMetadataRequest",
+            {
+                "doc_type": fields.String(description="Document type"),
+                "doc_metadata": fields.Raw(description="Document metadata"),
+            },
+        )
+    )
+    @api.response(200, "Document metadata updated successfully")
+    @api.response(404, "Document not found")
+    @api.response(403, "Permission denied")
     @setup_required
     @login_required
     @account_initialization_required