refactor: core/rag docstore, datasource, embedding, rerank, retrieval (#34203)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Asuka Minato <i@asukaminato.eu.org>

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

@@ -97,13 +97,13 @@ class Jieba(BaseKeyword):
 
         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 @@ class RetrievalService:
             # 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 @@ class TidbOnQdrantVectorFactory(AbstractVectorFactory):
                     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 @@ class Vector:
         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

@@ -4,7 +4,7 @@ from collections.abc import Sequence
 from typing import Any
 
 from graphon.model_runtime.entities.model_entities import ModelType
-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 @@ class DatasetDocumentStore:
         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 @@ class DatasetDocumentStore:
                 )
                 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

@@ -6,6 +6,7 @@ from typing import Any, cast
 import numpy as np
 from graphon.model_runtime.entities.model_entities import ModelPropertyKey
 from graphon.model_runtime.model_providers.__base.text_embedding_model import TextEmbeddingModel
+from sqlalchemy import select
 from sqlalchemy.exc import IntegrityError
 
 from configs import dify_config
@@ -31,14 +32,14 @@ class CacheEmbedding(Embeddings):
         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 @@ class CacheEmbedding(Embeddings):
         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

@@ -4,6 +4,7 @@ import operator
 from typing import Any, cast
 
 import httpx
+from sqlalchemy import update
 
 from configs import dify_config
 from core.rag.extractor.extractor_base import BaseExtractor
@@ -346,9 +347,11 @@ class NotionExtractor(BaseExtractor):
         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

@@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, NotRequired, Optional
 from urllib.parse import unquote, urlparse
 
 import httpx
+from sqlalchemy import select
 from typing_extensions import TypedDict
 
 from configs import dify_config
@@ -200,7 +201,7 @@ class BaseIndexProcessor(ABC):
 
         # 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 +313,7 @@ class BaseIndexProcessor(ABC):
         """
         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

@@ -18,6 +18,7 @@ from graphon.model_runtime.entities.message_entities import (
     UserPromptMessage,
 )
 from graphon.model_runtime.entities.model_entities import ModelFeature, ModelType
+from sqlalchemy import select
 
 from core.app.file_access import DatabaseFileAccessController
 from core.app.llm import deduct_llm_quota
@@ -145,14 +146,12 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
         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)
@@ -537,11 +536,9 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
 
         # 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

@@ -6,6 +6,8 @@ import uuid
 from collections.abc import Mapping
 from typing import Any
 
+from sqlalchemy import delete, select
+
 from configs import dify_config
 from core.db.session_factory import session_factory
 from core.entities.knowledge_entities import PreviewDetail
@@ -177,17 +179,16 @@ class ParentChildIndexProcessor(BaseIndexProcessor):
                     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:
@@ -195,18 +196,22 @@ class ParentChildIndexProcessor(BaseIndexProcessor):
 
                 # 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(

api/core/rag/rerank/rerank_model.py

@@ -134,9 +134,7 @@ class RerankModelRunner(BaseRerankRunner):
             ):
                 if document.metadata.get("doc_type") == DocType.IMAGE:
                     # Query file info within db.session context to ensure thread-safe access
-                    upload_file = (
-                        db.session.query(UploadFile).where(UploadFile.id == document.metadata["doc_id"]).first()
-                    )
+                    upload_file = db.session.get(UploadFile, document.metadata["doc_id"])
                     if upload_file:
                         blob = storage.load_once(upload_file.key)
                         document_file_base64 = base64.b64encode(blob).decode()
@@ -169,7 +167,7 @@ class RerankModelRunner(BaseRerankRunner):
             return rerank_result, unique_documents
         elif query_type == QueryType.IMAGE_QUERY:
             # Query file info within db.session context to ensure thread-safe access
-            upload_file = db.session.query(UploadFile).where(UploadFile.id == query).first()
+            upload_file = db.session.get(UploadFile, query)
             if upload_file:
                 blob = storage.load_once(upload_file.key)
                 file_query = base64.b64encode(blob).decode()

api/core/rag/retrieval/dataset_retrieval.py

@@ -1340,7 +1340,7 @@ class DatasetRetrieval:
         metadata_filtering_conditions: MetadataFilteringCondition | None,
         inputs: dict,
     ) -> tuple[dict[str, list[str]] | None, MetadataCondition | None]:
-        document_query = db.session.query(DatasetDocument).where(
+        document_query = select(DatasetDocument).where(
             DatasetDocument.dataset_id.in_(dataset_ids),
             DatasetDocument.indexing_status == "completed",
             DatasetDocument.enabled == True,
@@ -1411,7 +1411,7 @@ class DatasetRetrieval:
                 document_query = document_query.where(and_(*filters))
             else:
                 document_query = document_query.where(or_(*filters))
-        documents = document_query.all()
+        documents = db.session.scalars(document_query).all()
         # group by dataset_id
         metadata_filter_document_ids = defaultdict(list) if documents else None  # type: ignore
         for document in documents: