refactor: core/rag docstore, datasource, embedding, rerank, retrieval index processors

Author: RenzoMXD

api/core/rag/datasource/keyword/jieba/jieba.py

@@ -97,13 +97,13 @@ def search(self, query: str, **kwargs: Any) -> list[Document]:
 
         documents = []
 
-        segment_query_stmt = db.session.query(DocumentSegment).where(
+        segment_query_stmt = select(DocumentSegment).where(
             DocumentSegment.dataset_id == self.dataset.id, DocumentSegment.index_node_id.in_(sorted_chunk_indices)
         )
         if document_ids_filter:
             segment_query_stmt = segment_query_stmt.where(DocumentSegment.document_id.in_(document_ids_filter))
 
-        segments = db.session.execute(segment_query_stmt).scalars().all()
+        segments = db.session.scalars(segment_query_stmt).all()
         segment_map = {segment.index_node_id: segment for segment in segments}
         for chunk_index in sorted_chunk_indices:
             segment = segment_map.get(chunk_index)

api/core/rag/datasource/retrieval_service.py

@@ -432,10 +432,11 @@ def format_retrieval_documents(cls, documents: list[Document]) -> list[Retrieval
             # Batch query dataset documents
             dataset_documents = {
                 doc.id: doc
-                for doc in db.session.query(DatasetDocument)
-                .where(DatasetDocument.id.in_(document_ids))
-                .options(load_only(DatasetDocument.id, DatasetDocument.doc_form, DatasetDocument.dataset_id))
-                .all()
+                for doc in db.session.scalars(
+                    select(DatasetDocument)
+                    .where(DatasetDocument.id.in_(document_ids))
+                    .options(load_only(DatasetDocument.id, DatasetDocument.doc_form, DatasetDocument.dataset_id))
+                ).all()
             }
 
             valid_dataset_documents = {}

api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_on_qdrant_vector.py

@@ -426,11 +426,10 @@ def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings
                     TIDB_ON_QDRANT_API_KEY = f"{tidb_auth_binding.account}:{tidb_auth_binding.password}"
 
                 else:
-                    idle_tidb_auth_binding = (
-                        db.session.query(TidbAuthBinding)
+                    idle_tidb_auth_binding = db.session.scalar(
+                        select(TidbAuthBinding)
                         .where(TidbAuthBinding.active == False, TidbAuthBinding.status == "ACTIVE")
                         .limit(1)
-                        .one_or_none()
                     )
                     if idle_tidb_auth_binding:
                         idle_tidb_auth_binding.active = True

api/core/rag/datasource/vdb/vector_factory.py

@@ -277,7 +277,7 @@ def search_by_vector(self, query: str, **kwargs: Any) -> list[Document]:
         return self._vector_processor.search_by_vector(query_vector, **kwargs)
 
     def search_by_file(self, file_id: str, **kwargs: Any) -> list[Document]:
-        upload_file: UploadFile | None = db.session.query(UploadFile).where(UploadFile.id == file_id).first()
+        upload_file: UploadFile | None = db.session.get(UploadFile, file_id)
 
         if not upload_file:
             return []

api/core/rag/docstore/dataset_docstore.py

@@ -3,7 +3,7 @@
 from collections.abc import Sequence
 from typing import Any
 
-from sqlalchemy import func, select
+from sqlalchemy import delete, func, select
 
 from core.model_manager import ModelManager
 from core.rag.index_processor.constant.index_type import IndexTechniqueType
@@ -63,10 +63,8 @@ def docs(self) -> dict[str, Document]:
         return output
 
     def add_documents(self, docs: Sequence[Document], allow_update: bool = True, save_child: bool = False):
-        max_position = (
-            db.session.query(func.max(DocumentSegment.position))
-            .where(DocumentSegment.document_id == self._document_id)
-            .scalar()
+        max_position = db.session.scalar(
+            select(func.max(DocumentSegment.position)).where(DocumentSegment.document_id == self._document_id)
         )
 
         if max_position is None:
@@ -155,12 +153,14 @@ def add_documents(self, docs: Sequence[Document], allow_update: bool = True, sav
                 )
                 if save_child and doc.children:
                     # delete the existing child chunks
-                    db.session.query(ChildChunk).where(
-                        ChildChunk.tenant_id == self._dataset.tenant_id,
-                        ChildChunk.dataset_id == self._dataset.id,
-                        ChildChunk.document_id == self._document_id,
-                        ChildChunk.segment_id == segment_document.id,
-                    ).delete()
+                    db.session.execute(
+                        delete(ChildChunk).where(
+                            ChildChunk.tenant_id == self._dataset.tenant_id,
+                            ChildChunk.dataset_id == self._dataset.id,
+                            ChildChunk.document_id == self._document_id,
+                            ChildChunk.segment_id == segment_document.id,
+                        )
+                    )
                     # add new child chunks
                     for position, child in enumerate(doc.children, start=1):
                         child_segment = ChildChunk(

api/core/rag/embedding/cached_embedding.py

@@ -4,6 +4,7 @@
 from typing import Any, cast
 
 import numpy as np
+from sqlalchemy import select
 from sqlalchemy.exc import IntegrityError
 
 from configs import dify_config
@@ -31,14 +32,14 @@ def embed_documents(self, texts: list[str]) -> list[list[float]]:
         embedding_queue_indices = []
         for i, text in enumerate(texts):
             hash = helper.generate_text_hash(text)
-            embedding = (
-                db.session.query(Embedding)
-                .filter_by(
-                    model_name=self._model_instance.model_name,
-                    hash=hash,
-                    provider_name=self._model_instance.provider,
+            embedding = db.session.scalar(
+                select(Embedding)
+                .where(
+                    Embedding.model_name == self._model_instance.model_name,
+                    Embedding.hash == hash,
+                    Embedding.provider_name == self._model_instance.provider,
                 )
-                .first()
+                .limit(1)
             )
             if embedding:
                 text_embeddings[i] = embedding.get_embedding()
@@ -112,14 +113,14 @@ def embed_multimodal_documents(self, multimodel_documents: list[dict]) -> list[l
         embedding_queue_indices = []
         for i, multimodel_document in enumerate(multimodel_documents):
             file_id = multimodel_document["file_id"]
-            embedding = (
-                db.session.query(Embedding)
-                .filter_by(
-                    model_name=self._model_instance.model_name,
-                    hash=file_id,
-                    provider_name=self._model_instance.provider,
+            embedding = db.session.scalar(
+                select(Embedding)
+                .where(
+                    Embedding.model_name == self._model_instance.model_name,
+                    Embedding.hash == file_id,
+                    Embedding.provider_name == self._model_instance.provider,
                 )
-                .first()
+                .limit(1)
             )
             if embedding:
                 multimodel_embeddings[i] = embedding.get_embedding()

api/core/rag/extractor/notion_extractor.py

@@ -8,6 +8,8 @@
 from configs import dify_config
 from core.rag.extractor.extractor_base import BaseExtractor
 from core.rag.models.document import Document
+from sqlalchemy import update
+
 from extensions.ext_database import db
 from models.dataset import Document as DocumentModel
 from services.datasource_provider_service import DatasourceProviderService
@@ -346,9 +348,11 @@ def update_last_edited_time(self, document_model: DocumentModel | None):
         if data_source_info:
             data_source_info["last_edited_time"] = last_edited_time
 
-        db.session.query(DocumentModel).filter_by(id=document_model.id).update(
-            {DocumentModel.data_source_info: json.dumps(data_source_info)}
-        )  # type: ignore
+        db.session.execute(
+            update(DocumentModel)
+            .where(DocumentModel.id == document_model.id)
+            .values(data_source_info=json.dumps(data_source_info))
+        )
         db.session.commit()
 
     def get_notion_last_edited_time(self) -> str:

api/core/rag/index_processor/index_processor_base.py

@@ -26,6 +26,8 @@
     FixedRecursiveCharacterTextSplitter,
 )
 from core.rag.splitter.text_splitter import TextSplitter
+from sqlalchemy import select
+
 from extensions.ext_database import db
 from extensions.ext_storage import storage
 from models import Account, ToolFile
@@ -200,7 +202,7 @@ def _get_content_files(self, document: Document, current_user: Account | None =
 
         # Get unique IDs for database query
         unique_upload_file_ids = list(set(upload_file_id_list))
-        upload_files = db.session.query(UploadFile).where(UploadFile.id.in_(unique_upload_file_ids)).all()
+        upload_files = db.session.scalars(select(UploadFile).where(UploadFile.id.in_(unique_upload_file_ids))).all()
 
         # Create a mapping from ID to UploadFile for quick lookup
         upload_file_map = {upload_file.id: upload_file for upload_file in upload_files}
@@ -312,7 +314,7 @@ def _download_tool_file(self, tool_file_id: str, current_user: Account) -> str |
         """
         from services.file_service import FileService
 
-        tool_file = db.session.query(ToolFile).where(ToolFile.id == tool_file_id).first()
+        tool_file = db.session.get(ToolFile, tool_file_id)
         if not tool_file:
             return None
         blob = storage.load_once(tool_file.file_key)

api/core/rag/index_processor/processor/paragraph_index_processor.py

@@ -29,6 +29,8 @@
 from core.rag.retrieval.retrieval_methods import RetrievalMethod
 from core.tools.utils.text_processing_utils import remove_leading_symbols
 from core.workflow.file_reference import build_file_reference
+from sqlalchemy import select
+
 from extensions.ext_database import db
 from factories.file_factory import build_from_mapping
 from graphon.file import File, FileTransferMethod, FileType, file_manager
@@ -144,14 +146,12 @@ def clean(self, dataset: Dataset, node_ids: list[str] | None, with_keywords: boo
         if delete_summaries:
             if node_ids:
                 # Find segments by index_node_id
-                segments = (
-                    db.session.query(DocumentSegment)
-                    .filter(
+                segments = db.session.scalars(
+                    select(DocumentSegment).where(
                         DocumentSegment.dataset_id == dataset.id,
                         DocumentSegment.index_node_id.in_(node_ids),
                     )
-                    .all()
-                )
+                ).all()
                 segment_ids = [segment.id for segment in segments]
                 if segment_ids:
                     SummaryIndexService.delete_summaries_for_segments(dataset, segment_ids)
@@ -536,11 +536,9 @@ def _extract_images_from_text(tenant_id: str, text: str) -> list[File]:
 
         # Get unique IDs for database query
         unique_upload_file_ids = list(set(upload_file_id_list))
-        upload_files = (
-            db.session.query(UploadFile)
-            .where(UploadFile.id.in_(unique_upload_file_ids), UploadFile.tenant_id == tenant_id)
-            .all()
-        )
+        upload_files = db.session.scalars(
+            select(UploadFile).where(UploadFile.id.in_(unique_upload_file_ids), UploadFile.tenant_id == tenant_id)
+        ).all()
 
         # Create File objects from UploadFile records
         file_objects = []

api/core/rag/index_processor/processor/parent_child_index_processor.py

@@ -22,6 +22,8 @@
 from core.rag.index_processor.index_processor_base import BaseIndexProcessor, SummaryIndexSettingDict
 from core.rag.models.document import AttachmentDocument, ChildDocument, Document, ParentChildStructureChunk
 from core.rag.retrieval.retrieval_methods import RetrievalMethod
+from sqlalchemy import delete, select
+
 from extensions.ext_database import db
 from libs import helper
 from models import Account
@@ -177,36 +179,39 @@ def clean(self, dataset: Dataset, node_ids: list[str] | None, with_keywords: boo
                     child_node_ids = precomputed_child_node_ids
                 else:
                     # Fallback to original query (may fail if segments are already deleted)
-                    child_node_ids = (
-                        db.session.query(ChildChunk.index_node_id)
+                    rows = db.session.execute(
+                        select(ChildChunk.index_node_id)
                         .join(DocumentSegment, ChildChunk.segment_id == DocumentSegment.id)
                         .where(
                             DocumentSegment.dataset_id == dataset.id,
                             DocumentSegment.index_node_id.in_(node_ids),
                             ChildChunk.dataset_id == dataset.id,
                         )
-                        .all()
-                    )
-                    child_node_ids = [child_node_id[0] for child_node_id in child_node_ids if child_node_id[0]]
+                    ).all()
+                    child_node_ids = [row[0] for row in rows if row[0]]
 
                 # Delete from vector index
                 if child_node_ids:
                     vector.delete_by_ids(child_node_ids)
 
                 # Delete from database
                 if delete_child_chunks and child_node_ids:
-                    db.session.query(ChildChunk).where(
-                        ChildChunk.dataset_id == dataset.id, ChildChunk.index_node_id.in_(child_node_ids)
-                    ).delete(synchronize_session=False)
+                    db.session.execute(
+                        delete(ChildChunk).where(
+                            ChildChunk.dataset_id == dataset.id, ChildChunk.index_node_id.in_(child_node_ids)
+                        )
+                    )
                     db.session.commit()
             else:
                 vector.delete()
 
                 if delete_child_chunks:
                     # Use existing compound index: (tenant_id, dataset_id, ...)
-                    db.session.query(ChildChunk).where(
-                        ChildChunk.tenant_id == dataset.tenant_id, ChildChunk.dataset_id == dataset.id
-                    ).delete(synchronize_session=False)
+                    db.session.execute(
+                        delete(ChildChunk).where(
+                            ChildChunk.tenant_id == dataset.tenant_id, ChildChunk.dataset_id == dataset.id
+                        )
+                    )
                     db.session.commit()
 
     def retrieve(