Merge branch 'main' into feat/rag-2

api/controllers/console/datasets/data_source.py

@@ -4,7 +4,7 @@ from typing import cast
 
 from flask import request
 from flask_login import current_user
-from flask_restful import Resource, marshal_with, reqparse
+from flask_restx import Resource, marshal_with, reqparse
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 from werkzeug.exceptions import NotFound

api/controllers/console/datasets/datasets.py

@@ -1,7 +1,7 @@
-import flask_restful
+import flask_restx
 from flask import request
 from flask_login import current_user
-from flask_restful import Resource, marshal, marshal_with, reqparse
+from flask_restx import Resource, marshal, marshal_with, reqparse
 from werkzeug.exceptions import Forbidden, NotFound
 
 import services
@@ -600,7 +600,7 @@ class DatasetApiKeyApi(Resource):
         )
 
         if current_key_count >= self.max_keys:
-            flask_restful.abort(
+            flask_restx.abort(
                 400,
                 message=f"Cannot create more than {self.max_keys} API keys for this resource type.",
                 code="max_keys_exceeded",
@@ -640,7 +640,7 @@ class DatasetApiDeleteApi(Resource):
         )
 
         if key is None:
-            flask_restful.abort(404, message="API key not found")
+            flask_restx.abort(404, message="API key not found")
 
         db.session.query(ApiToken).where(ApiToken.id == api_key_id).delete()
         db.session.commit()

api/controllers/console/datasets/datasets_document.py

@@ -5,7 +5,7 @@ from typing import Literal, cast
 
 from flask import request
 from flask_login import current_user
-from flask_restful import Resource, marshal, marshal_with, reqparse
+from flask_restx import Resource, marshal, marshal_with, reqparse
 from sqlalchemy import asc, desc, select
 from werkzeug.exceptions import Forbidden, NotFound
 

api/controllers/console/datasets/datasets_segments.py

@@ -2,7 +2,7 @@ import uuid
 
 from flask import request
 from flask_login import current_user
-from flask_restful import Resource, marshal, reqparse
+from flask_restx import Resource, marshal, reqparse
 from sqlalchemy import select
 from werkzeug.exceptions import Forbidden, NotFound
 
@@ -584,7 +584,12 @@ class ChildChunkUpdateApi(Resource):
         child_chunk_id = str(child_chunk_id)
         child_chunk = (
             db.session.query(ChildChunk)
-            .where(ChildChunk.id == str(child_chunk_id), ChildChunk.tenant_id == current_user.current_tenant_id)
+            .where(
+                ChildChunk.id == str(child_chunk_id),
+                ChildChunk.tenant_id == current_user.current_tenant_id,
+                ChildChunk.segment_id == segment.id,
+                ChildChunk.document_id == document_id,
+            )
             .first()
         )
         if not child_chunk:
@@ -633,7 +638,12 @@ class ChildChunkUpdateApi(Resource):
         child_chunk_id = str(child_chunk_id)
         child_chunk = (
             db.session.query(ChildChunk)
-            .where(ChildChunk.id == str(child_chunk_id), ChildChunk.tenant_id == current_user.current_tenant_id)
+            .where(
+                ChildChunk.id == str(child_chunk_id),
+                ChildChunk.tenant_id == current_user.current_tenant_id,
+                ChildChunk.segment_id == segment.id,
+                ChildChunk.document_id == document_id,
+            )
             .first()
         )
         if not child_chunk:

api/controllers/console/datasets/external.py

@@ -1,6 +1,6 @@
 from flask import request
 from flask_login import current_user
-from flask_restful import Resource, marshal, reqparse
+from flask_restx import Resource, marshal, reqparse
 from werkzeug.exceptions import Forbidden, InternalServerError, NotFound
 
 import services

api/controllers/console/datasets/hit_testing.py

@@ -1,4 +1,4 @@
-from flask_restful import Resource
+from flask_restx import Resource
 
 from controllers.console import api
 from controllers.console.datasets.hit_testing_base import DatasetsHitTestingBase

api/controllers/console/datasets/hit_testing_base.py

@@ -1,7 +1,7 @@
 import logging
 
 from flask_login import current_user
-from flask_restful import marshal, reqparse
+from flask_restx import marshal, reqparse
 from werkzeug.exceptions import Forbidden, InternalServerError, NotFound
 
 import services.dataset_service

api/controllers/console/datasets/metadata.py

@@ -1,7 +1,7 @@
 from typing import Literal
 
 from flask_login import current_user
-from flask_restful import Resource, marshal_with, reqparse
+from flask_restx import Resource, marshal_with, reqparse
 from werkzeug.exceptions import NotFound
 
 from controllers.console import api

api/controllers/console/datasets/upload_file.py

@@ -1,62 +0,0 @@
-from flask_login import current_user
-from flask_restful import Resource
-from werkzeug.exceptions import NotFound
-
-from controllers.console import api
-from controllers.console.wraps import (
-    account_initialization_required,
-    setup_required,
-)
-from core.file import helpers as file_helpers
-from extensions.ext_database import db
-from models.dataset import Dataset
-from models.model import UploadFile
-from services.dataset_service import DocumentService
-
-
-class UploadFileApi(Resource):
-    @setup_required
-    @account_initialization_required
-    def get(self, dataset_id, document_id):
-        """Get upload file."""
-        # check dataset
-        dataset_id = str(dataset_id)
-        dataset = (
-            db.session.query(Dataset)
-            .filter(Dataset.tenant_id == current_user.current_tenant_id, Dataset.id == dataset_id)
-            .first()
-        )
-        if not dataset:
-            raise NotFound("Dataset not found.")
-        # check document
-        document_id = str(document_id)
-        document = DocumentService.get_document(dataset.id, document_id)
-        if not document:
-            raise NotFound("Document not found.")
-        # check upload file
-        if document.data_source_type != "upload_file":
-            raise ValueError(f"Document data source type ({document.data_source_type}) is not upload_file.")
-        data_source_info = document.data_source_info_dict
-        if data_source_info and "upload_file_id" in data_source_info:
-            file_id = data_source_info["upload_file_id"]
-            upload_file = db.session.query(UploadFile).where(UploadFile.id == file_id).first()
-            if not upload_file:
-                raise NotFound("UploadFile not found.")
-        else:
-            raise ValueError("Upload file id not found in document data source info.")
-
-        url = file_helpers.get_signed_file_url(upload_file_id=upload_file.id)
-        return {
-            "id": upload_file.id,
-            "name": upload_file.name,
-            "size": upload_file.size,
-            "extension": upload_file.extension,
-            "url": url,
-            "download_url": f"{url}&as_attachment=true",
-            "mime_type": upload_file.mime_type,
-            "created_by": upload_file.created_by,
-            "created_at": upload_file.created_at.timestamp(),
-        }, 200
-
-
-api.add_resource(UploadFileApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/upload-file")

api/controllers/console/datasets/website.py

@@ -1,4 +1,4 @@
-from flask_restful import Resource, reqparse
+from flask_restx import Resource, reqparse
 
 from controllers.console import api
 from controllers.console.datasets.error import WebsiteCrawlError