Merge remote-tracking branch 'origin/main' into feat/queue-based-graph-engine

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

@@ -3,6 +3,7 @@ from typing import Any, Optional
 
 import orjson
 from pydantic import BaseModel
+from sqlalchemy import select
 
 from configs import dify_config
 from core.rag.datasource.keyword.jieba.jieba_keyword_table_handler import JiebaKeywordTableHandler
@@ -211,11 +212,10 @@ class Jieba(BaseKeyword):
         return sorted_chunk_indices[:k]
 
     def _update_segment_keywords(self, dataset_id: str, node_id: str, keywords: list[str]):
-        document_segment = (
-            db.session.query(DocumentSegment)
-            .where(DocumentSegment.dataset_id == dataset_id, DocumentSegment.index_node_id == node_id)
-            .first()
+        stmt = select(DocumentSegment).where(
+            DocumentSegment.dataset_id == dataset_id, DocumentSegment.index_node_id == node_id
         )
+        document_segment = db.session.scalar(stmt)
         if document_segment:
             document_segment.keywords = keywords
             db.session.add(document_segment)

api/core/rag/datasource/retrieval_service.py

@@ -3,6 +3,7 @@ from concurrent.futures import ThreadPoolExecutor
 from typing import Optional
 
 from flask import Flask, current_app
+from sqlalchemy import select
 from sqlalchemy.orm import Session, load_only
 
 from configs import dify_config
@@ -24,7 +25,7 @@ default_retrieval_model = {
     "search_method": RetrievalMethod.SEMANTIC_SEARCH.value,
     "reranking_enable": False,
     "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""},
-    "top_k": 2,
+    "top_k": 4,
     "score_threshold_enabled": False,
 }
 
@@ -127,7 +128,8 @@ class RetrievalService:
         external_retrieval_model: Optional[dict] = None,
         metadata_filtering_conditions: Optional[dict] = None,
     ):
-        dataset = db.session.query(Dataset).where(Dataset.id == dataset_id).first()
+        stmt = select(Dataset).where(Dataset.id == dataset_id)
+        dataset = db.session.scalar(stmt)
         if not dataset:
             return []
         metadata_condition = (
@@ -316,10 +318,8 @@ class RetrievalService:
                 if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX:
                     # Handle parent-child documents
                     child_index_node_id = document.metadata.get("doc_id")
-
-                    child_chunk = (
-                        db.session.query(ChildChunk).where(ChildChunk.index_node_id == child_index_node_id).first()
-                    )
+                    child_chunk_stmt = select(ChildChunk).where(ChildChunk.index_node_id == child_index_node_id)
+                    child_chunk = db.session.scalar(child_chunk_stmt)
 
                     if not child_chunk:
                         continue
@@ -378,17 +378,13 @@ class RetrievalService:
                     index_node_id = document.metadata.get("doc_id")
                     if not index_node_id:
                         continue
-
-                    segment = (
-                        db.session.query(DocumentSegment)
-                        .where(
-                            DocumentSegment.dataset_id == dataset_document.dataset_id,
-                            DocumentSegment.enabled == True,
-                            DocumentSegment.status == "completed",
-                            DocumentSegment.index_node_id == index_node_id,
-                        )
-                        .first()
+                    document_segment_stmt = select(DocumentSegment).where(
+                        DocumentSegment.dataset_id == dataset_document.dataset_id,
+                        DocumentSegment.enabled == True,
+                        DocumentSegment.status == "completed",
+                        DocumentSegment.index_node_id == index_node_id,
                     )
+                    segment = db.session.scalar(document_segment_stmt)
 
                     if not segment:
                         continue

api/core/rag/datasource/vdb/analyticdb/analyticdb_vector_openapi.py

@@ -256,7 +256,7 @@ class AnalyticdbVectorOpenAPI:
         response = self._client.query_collection_data(request)
         documents = []
         for match in response.body.matches.match:
-            if match.score > score_threshold:
+            if match.score >= score_threshold:
                 metadata = json.loads(match.metadata.get("metadata_"))
                 metadata["score"] = match.score
                 doc = Document(
@@ -293,7 +293,7 @@ class AnalyticdbVectorOpenAPI:
         response = self._client.query_collection_data(request)
         documents = []
         for match in response.body.matches.match:
-            if match.score > score_threshold:
+            if match.score >= score_threshold:
                 metadata = json.loads(match.metadata.get("metadata_"))
                 metadata["score"] = match.score
                 doc = Document(

api/core/rag/datasource/vdb/analyticdb/analyticdb_vector_sql.py

@@ -229,7 +229,7 @@ class AnalyticdbVectorBySql:
             documents = []
             for record in cur:
                 id, vector, score, page_content, metadata = record
-                if score > score_threshold:
+                if score >= score_threshold:
                     metadata["score"] = score
                     doc = Document(
                         page_content=page_content,

api/core/rag/datasource/vdb/baidu/baidu_vector.py

@@ -157,7 +157,7 @@ class BaiduVector(BaseVector):
             if meta is not None:
                 meta = json.loads(meta)
             score = row.get("score", 0.0)
-            if score > score_threshold:
+            if score >= score_threshold:
                 meta["score"] = score
                 doc = Document(page_content=row_data.get(self.field_text), metadata=meta)
                 docs.append(doc)

api/core/rag/datasource/vdb/chroma/chroma_vector.py

@@ -120,7 +120,7 @@ class ChromaVector(BaseVector):
             distance = distances[index]
             metadata = dict(metadatas[index])
             score = 1 - distance
-            if score > score_threshold:
+            if score >= score_threshold:
                 metadata["score"] = score
                 doc = Document(
                     page_content=documents[index],

api/core/rag/datasource/vdb/couchbase/couchbase_vector.py

@@ -304,7 +304,7 @@ class CouchbaseVector(BaseVector):
         return docs
 
     def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]:
-        top_k = kwargs.get("top_k", 2)
+        top_k = kwargs.get("top_k", 4)
         try:
             CBrequest = search.SearchRequest.create(search.QueryStringQuery("text:" + query))
             search_iter = self._scope.search(

api/core/rag/datasource/vdb/elasticsearch/elasticsearch_vector.py

@@ -216,7 +216,7 @@ class ElasticSearchVector(BaseVector):
         docs = []
         for doc, score in docs_and_scores:
             score_threshold = float(kwargs.get("score_threshold") or 0.0)
-            if score > score_threshold:
+            if score >= score_threshold:
                 if doc.metadata is not None:
                     doc.metadata["score"] = score
                     docs.append(doc)

api/core/rag/datasource/vdb/huawei/huawei_cloud_vector.py

@@ -127,7 +127,7 @@ class HuaweiCloudVector(BaseVector):
         docs = []
         for doc, score in docs_and_scores:
             score_threshold = float(kwargs.get("score_threshold") or 0.0)
-            if score > score_threshold:
+            if score >= score_threshold:
                 if doc.metadata is not None:
                     doc.metadata["score"] = score
             docs.append(doc)

api/core/rag/datasource/vdb/lindorm/lindorm_vector.py

@@ -275,7 +275,7 @@ class LindormVectorStore(BaseVector):
         docs = []
         for doc, score in docs_and_scores:
             score_threshold = kwargs.get("score_threshold", 0.0) or 0.0
-            if score > score_threshold:
+            if score >= score_threshold:
                 if doc.metadata is not None:
                     doc.metadata["score"] = score
                 docs.append(doc)

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

@@ -194,7 +194,7 @@ class OpenGauss(BaseVector):
                 metadata, text, distance = record
                 score = 1 - distance
                 metadata["score"] = score
-                if score > score_threshold:
+                if score >= score_threshold:
                     docs.append(Document(page_content=text, metadata=metadata))
         return docs
 

api/core/rag/datasource/vdb/opensearch/opensearch_vector.py

@@ -211,7 +211,7 @@ class OpenSearchVector(BaseVector):
 
             metadata["score"] = hit["_score"]
             score_threshold = float(kwargs.get("score_threshold") or 0.0)
-            if hit["_score"] > score_threshold:
+            if hit["_score"] >= score_threshold:
                 doc = Document(page_content=hit["_source"].get(Field.CONTENT_KEY.value), metadata=metadata)
                 docs.append(doc)
 

api/core/rag/datasource/vdb/oracle/oraclevector.py

@@ -261,7 +261,7 @@ class OracleVector(BaseVector):
                     metadata, text, distance = record
                     score = 1 - distance
                     metadata["score"] = score
-                    if score > score_threshold:
+                    if score >= score_threshold:
                         docs.append(Document(page_content=text, metadata=metadata))
             conn.close()
         return docs

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

@@ -202,7 +202,7 @@ class PGVectoRS(BaseVector):
             score = 1 - dis
             metadata["score"] = score
             score_threshold = float(kwargs.get("score_threshold") or 0.0)
-            if score > score_threshold:
+            if score >= score_threshold:
                 doc = Document(page_content=record.text, metadata=metadata)
                 docs.append(doc)
         return docs

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

@@ -195,7 +195,7 @@ class PGVector(BaseVector):
                 metadata, text, distance = record
                 score = 1 - distance
                 metadata["score"] = score
-                if score > score_threshold:
+                if score >= score_threshold:
                     docs.append(Document(page_content=text, metadata=metadata))
         return docs
 

api/core/rag/datasource/vdb/pyvastbase/vastbase_vector.py

@@ -170,7 +170,7 @@ class VastbaseVector(BaseVector):
                 metadata, text, distance = record
                 score = 1 - distance
                 metadata["score"] = score
-                if score > score_threshold:
+                if score >= score_threshold:
                     docs.append(Document(page_content=text, metadata=metadata))
         return docs
 

api/core/rag/datasource/vdb/qdrant/qdrant_vector.py

@@ -3,7 +3,7 @@ import os
 import uuid
 from collections.abc import Generator, Iterable, Sequence
 from itertools import islice
-from typing import TYPE_CHECKING, Any, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 import qdrant_client
 from flask import current_app
@@ -18,6 +18,7 @@ from qdrant_client.http.models import (
     TokenizerType,
 )
 from qdrant_client.local.qdrant_local import QdrantLocal
+from sqlalchemy import select
 
 from configs import dify_config
 from core.rag.datasource.vdb.field import Field
@@ -369,7 +370,7 @@ class QdrantVector(BaseVector):
                 continue
             metadata = result.payload.get(Field.METADATA_KEY.value) or {}
             # duplicate check score threshold
-            if result.score > score_threshold:
+            if result.score >= score_threshold:
                 metadata["score"] = result.score
                 doc = Document(
                     page_content=result.payload.get(Field.CONTENT_KEY.value, ""),
@@ -426,7 +427,6 @@ class QdrantVector(BaseVector):
 
     def _reload_if_needed(self):
         if isinstance(self._client, QdrantLocal):
-            self._client = cast(QdrantLocal, self._client)
             self._client._load()
 
     @classmethod
@@ -446,11 +446,8 @@ class QdrantVector(BaseVector):
 class QdrantVectorFactory(AbstractVectorFactory):
     def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> QdrantVector:
         if dataset.collection_binding_id:
-            dataset_collection_binding = (
-                db.session.query(DatasetCollectionBinding)
-                .where(DatasetCollectionBinding.id == dataset.collection_binding_id)
-                .one_or_none()
-            )
+            stmt = select(DatasetCollectionBinding).where(DatasetCollectionBinding.id == dataset.collection_binding_id)
+            dataset_collection_binding = db.session.scalars(stmt).one_or_none()
             if dataset_collection_binding:
                 collection_name = dataset_collection_binding.collection_name
             else:

api/core/rag/datasource/vdb/relyt/relyt_vector.py

@@ -233,7 +233,7 @@ class RelytVector(BaseVector):
         docs = []
         for document, score in results:
             score_threshold = float(kwargs.get("score_threshold") or 0.0)
-            if 1 - score > score_threshold:
+            if 1 - score >= score_threshold:
                 docs.append(document)
         return docs
 

api/core/rag/datasource/vdb/tablestore/tablestore_vector.py

@@ -300,7 +300,7 @@ class TableStoreVector(BaseVector):
         )
         documents = []
         for search_hit in search_response.search_hits:
-            if search_hit.score > score_threshold:
+            if search_hit.score >= score_threshold:
                 ots_column_map = {}
                 for col in search_hit.row[1]:
                     ots_column_map[col[0]] = col[1]

api/core/rag/datasource/vdb/tencent/tencent_vector.py

@@ -39,6 +39,9 @@ class TencentConfig(BaseModel):
         return {"url": self.url, "username": self.username, "key": self.api_key, "timeout": self.timeout}
 
 
+bm25 = BM25Encoder.default("zh")
+
+
 class TencentVector(BaseVector):
     field_id: str = "id"
     field_vector: str = "vector"
@@ -53,7 +56,6 @@ class TencentVector(BaseVector):
         self._dimension = 1024
         self._init_database()
         self._load_collection()
-        self._bm25 = BM25Encoder.default("zh")
 
     def _load_collection(self):
         """
@@ -186,7 +188,7 @@ class TencentVector(BaseVector):
                     metadata=metadata,
                 )
                 if self._enable_hybrid_search:
-                    doc.__dict__["sparse_vector"] = self._bm25.encode_texts(texts[i])
+                    doc.__dict__["sparse_vector"] = bm25.encode_texts(texts[i])
                 docs.append(doc)
             self._client.upsert(
                 database_name=self._client_config.database,
@@ -264,7 +266,7 @@ class TencentVector(BaseVector):
             match=[
                 KeywordSearch(
                     field_name="sparse_vector",
-                    data=self._bm25.encode_queries(query),
+                    data=bm25.encode_queries(query),
                 ),
             ],
             rerank=WeightedRerank(
@@ -291,7 +293,7 @@ class TencentVector(BaseVector):
                 score = 1 - result.get("score", 0.0)
             else:
                 score = result.get("score", 0.0)
-            if score > score_threshold:
+            if score >= score_threshold:
                 meta["score"] = score
                 doc = Document(page_content=result.get(self.field_text), metadata=meta)
                 docs.append(doc)

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

@@ -20,6 +20,7 @@ from qdrant_client.http.models import (
 )
 from qdrant_client.local.qdrant_local import QdrantLocal
 from requests.auth import HTTPDigestAuth
+from sqlalchemy import select
 
 from configs import dify_config
 from core.rag.datasource.vdb.field import Field
@@ -351,7 +352,7 @@ class TidbOnQdrantVector(BaseVector):
             metadata = result.payload.get(Field.METADATA_KEY.value) or {}
             # duplicate check score threshold
             score_threshold = kwargs.get("score_threshold") or 0.0
-            if result.score > score_threshold:
+            if result.score >= score_threshold:
                 metadata["score"] = result.score
                 doc = Document(
                     page_content=result.payload.get(Field.CONTENT_KEY.value, ""),
@@ -416,16 +417,12 @@ class TidbOnQdrantVector(BaseVector):
 
 class TidbOnQdrantVectorFactory(AbstractVectorFactory):
     def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> TidbOnQdrantVector:
-        tidb_auth_binding = (
-            db.session.query(TidbAuthBinding).where(TidbAuthBinding.tenant_id == dataset.tenant_id).one_or_none()
-        )
+        stmt = select(TidbAuthBinding).where(TidbAuthBinding.tenant_id == dataset.tenant_id)
+        tidb_auth_binding = db.session.scalars(stmt).one_or_none()
         if not tidb_auth_binding:
             with redis_client.lock("create_tidb_serverless_cluster_lock", timeout=900):
-                tidb_auth_binding = (
-                    db.session.query(TidbAuthBinding)
-                    .where(TidbAuthBinding.tenant_id == dataset.tenant_id)
-                    .one_or_none()
-                )
+                stmt = select(TidbAuthBinding).where(TidbAuthBinding.tenant_id == dataset.tenant_id)
+                tidb_auth_binding = db.session.scalars(stmt).one_or_none()
                 if tidb_auth_binding:
                     TIDB_ON_QDRANT_API_KEY = f"{tidb_auth_binding.account}:{tidb_auth_binding.password}"
 

api/core/rag/datasource/vdb/upstash/upstash_vector.py

@@ -110,7 +110,7 @@ class UpstashVector(BaseVector):
             score = record.score
             if metadata is not None and text is not None:
                 metadata["score"] = score
-                if score > score_threshold:
+                if score >= score_threshold:
                     docs.append(Document(page_content=text, metadata=metadata))
         return docs
 

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

@@ -3,6 +3,8 @@ import time
 from abc import ABC, abstractmethod
 from typing import Any, Optional
 
+from sqlalchemy import select
+
 from configs import dify_config
 from core.model_manager import ModelManager
 from core.model_runtime.entities.model_entities import ModelType
@@ -45,11 +47,10 @@ class Vector:
             vector_type = self._dataset.index_struct_dict["type"]
         else:
             if dify_config.VECTOR_STORE_WHITELIST_ENABLE:
-                whitelist = (
-                    db.session.query(Whitelist)
-                    .where(Whitelist.tenant_id == self._dataset.tenant_id, Whitelist.category == "vector_db")
-                    .one_or_none()
+                stmt = select(Whitelist).where(
+                    Whitelist.tenant_id == self._dataset.tenant_id, Whitelist.category == "vector_db"
                 )
+                whitelist = db.session.scalars(stmt).one_or_none()
                 if whitelist:
                     vector_type = VectorType.TIDB_ON_QDRANT
 

api/core/rag/datasource/vdb/vikingdb/vikingdb_vector.py

@@ -192,7 +192,7 @@ class VikingDBVector(BaseVector):
             metadata = result.fields.get(vdb_Field.METADATA_KEY.value)
             if metadata is not None:
                 metadata = json.loads(metadata)
-            if result.score > score_threshold:
+            if result.score >= score_threshold:
                 metadata["score"] = result.score
                 doc = Document(page_content=result.fields.get(vdb_Field.CONTENT_KEY.value), metadata=metadata)
                 docs.append(doc)

api/core/rag/datasource/vdb/weaviate/weaviate_vector.py

@@ -220,7 +220,7 @@ class WeaviateVector(BaseVector):
         for doc, score in docs_and_scores:
             score_threshold = float(kwargs.get("score_threshold") or 0.0)
             # check score threshold
-            if score > score_threshold:
+            if score >= score_threshold:
                 if doc.metadata is not None:
                     doc.metadata["score"] = score
                     docs.append(doc)

api/core/rag/docstore/dataset_docstore.py

@@ -1,7 +1,7 @@
 from collections.abc import Sequence
 from typing import Any, Optional
 
-from sqlalchemy import func
+from sqlalchemy import func, select
 
 from core.model_manager import ModelManager
 from core.model_runtime.entities.model_entities import ModelType
@@ -41,9 +41,8 @@ class DatasetDocumentStore:
 
     @property
     def docs(self) -> dict[str, Document]:
-        document_segments = (
-            db.session.query(DocumentSegment).where(DocumentSegment.dataset_id == self._dataset.id).all()
-        )
+        stmt = select(DocumentSegment).where(DocumentSegment.dataset_id == self._dataset.id)
+        document_segments = db.session.scalars(stmt).all()
 
         output = {}
         for document_segment in document_segments:
@@ -228,10 +227,9 @@ class DatasetDocumentStore:
         return data
 
     def get_document_segment(self, doc_id: str) -> Optional[DocumentSegment]:
-        document_segment = (
-            db.session.query(DocumentSegment)
-            .where(DocumentSegment.dataset_id == self._dataset.id, DocumentSegment.index_node_id == doc_id)
-            .first()
+        stmt = select(DocumentSegment).where(
+            DocumentSegment.dataset_id == self._dataset.id, DocumentSegment.index_node_id == doc_id
         )
+        document_segment = db.session.scalar(stmt)
 
         return document_segment

api/core/rag/extractor/markdown_extractor.py

@@ -2,7 +2,7 @@
 
 import re
 from pathlib import Path
-from typing import Optional, cast
+from typing import Optional
 
 from core.rag.extractor.extractor_base import BaseExtractor
 from core.rag.extractor.helpers import detect_file_encodings
@@ -76,7 +76,7 @@ class MarkdownExtractor(BaseExtractor):
         markdown_tups.append((current_header, current_text))
 
         markdown_tups = [
-            (re.sub(r"#", "", cast(str, key)).strip() if key else None, re.sub(r"<.*?>", "", value))
+            (re.sub(r"#", "", key).strip() if key else None, re.sub(r"<.*?>", "", value))
             for key, value in markdown_tups
         ]
 

api/core/rag/extractor/notion_extractor.py

@@ -4,6 +4,7 @@ import operator
 from typing import Any, Optional, cast
 
 import requests
+from sqlalchemy import select
 
 from configs import dify_config
 from core.rag.extractor.extractor_base import BaseExtractor
@@ -367,22 +368,17 @@ class NotionExtractor(BaseExtractor):
 
     @classmethod
     def _get_access_token(cls, tenant_id: str, notion_workspace_id: str) -> str:
-        data_source_binding = (
-            db.session.query(DataSourceOauthBinding)
-            .where(
-                db.and_(
-                    DataSourceOauthBinding.tenant_id == tenant_id,
-                    DataSourceOauthBinding.provider == "notion",
-                    DataSourceOauthBinding.disabled == False,
-                    DataSourceOauthBinding.source_info["workspace_id"] == f'"{notion_workspace_id}"',
-                )
-            )
-            .first()
+        stmt = select(DataSourceOauthBinding).where(
+            DataSourceOauthBinding.tenant_id == tenant_id,
+            DataSourceOauthBinding.provider == "notion",
+            DataSourceOauthBinding.disabled == False,
+            DataSourceOauthBinding.source_info["workspace_id"] == f'"{notion_workspace_id}"',
         )
+        data_source_binding = db.session.scalar(stmt)
 
         if not data_source_binding:
             raise Exception(
                 f"No notion data source binding found for tenant {tenant_id} and notion workspace {notion_workspace_id}"
             )
 
-        return cast(str, data_source_binding.access_token)
+        return data_source_binding.access_token

api/core/rag/extractor/pdf_extractor.py

@@ -2,7 +2,7 @@
 
 import contextlib
 from collections.abc import Iterator
-from typing import Optional, cast
+from typing import Optional
 
 from core.rag.extractor.blob.blob import Blob
 from core.rag.extractor.extractor_base import BaseExtractor
@@ -27,7 +27,7 @@ class PdfExtractor(BaseExtractor):
         plaintext_file_exists = False
         if self._file_cache_key:
             with contextlib.suppress(FileNotFoundError):
-                text = cast(bytes, storage.load(self._file_cache_key)).decode("utf-8")
+                text = storage.load(self._file_cache_key).decode("utf-8")
                 plaintext_file_exists = True
                 return [Document(page_content=text)]
         documents = list(self.load())

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

@@ -123,7 +123,7 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
         for result in results:
             metadata = result.metadata
             metadata["score"] = result.score
-            if result.score > score_threshold:
+            if result.score >= score_threshold:
                 doc = Document(page_content=result.page_content, metadata=metadata)
                 docs.append(doc)
         return docs

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

@@ -130,13 +130,16 @@ class ParentChildIndexProcessor(BaseIndexProcessor):
                 if delete_child_chunks:
                     db.session.query(ChildChunk).where(
                         ChildChunk.dataset_id == dataset.id, ChildChunk.index_node_id.in_(child_node_ids)
-                    ).delete()
+                    ).delete(synchronize_session=False)
                     db.session.commit()
             else:
                 vector.delete()
 
                 if delete_child_chunks:
-                    db.session.query(ChildChunk).where(ChildChunk.dataset_id == dataset.id).delete()
+                    # 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.commit()
 
     def retrieve(
@@ -162,7 +165,7 @@ class ParentChildIndexProcessor(BaseIndexProcessor):
         for result in results:
             metadata = result.metadata
             metadata["score"] = result.score
-            if result.score > score_threshold:
+            if result.score >= score_threshold:
                 doc = Document(page_content=result.page_content, metadata=metadata)
                 docs.append(doc)
         return docs

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

@@ -158,7 +158,7 @@ class QAIndexProcessor(BaseIndexProcessor):
         for result in results:
             metadata = result.metadata
             metadata["score"] = result.score
-            if result.score > score_threshold:
+            if result.score >= score_threshold:
                 doc = Document(page_content=result.page_content, metadata=metadata)
                 docs.append(doc)
         return docs

api/core/rag/retrieval/dataset_retrieval.py

@@ -7,7 +7,7 @@ from collections.abc import Generator, Mapping
 from typing import Any, Optional, Union, cast
 
 from flask import Flask, current_app
-from sqlalchemy import Float, and_, or_, text
+from sqlalchemy import Float, and_, or_, select, text
 from sqlalchemy import cast as sqlalchemy_cast
 from sqlalchemy.orm import Session
 
@@ -65,7 +65,7 @@ default_retrieval_model: dict[str, Any] = {
     "search_method": RetrievalMethod.SEMANTIC_SEARCH.value,
     "reranking_enable": False,
     "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""},
-    "top_k": 2,
+    "top_k": 4,
     "score_threshold_enabled": False,
 }
 
@@ -135,7 +135,8 @@ class DatasetRetrieval:
         available_datasets = []
         for dataset_id in dataset_ids:
             # get dataset from dataset id
-            dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first()
+            dataset_stmt = select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id)
+            dataset = db.session.scalar(dataset_stmt)
 
             # pass if dataset is not available
             if not dataset:
@@ -240,15 +241,12 @@ class DatasetRetrieval:
                     for record in records:
                         segment = record.segment
                         dataset = db.session.query(Dataset).filter_by(id=segment.dataset_id).first()
-                        document = (
-                            db.session.query(DatasetDocument)
-                            .where(
-                                DatasetDocument.id == segment.document_id,
-                                DatasetDocument.enabled == True,
-                                DatasetDocument.archived == False,
-                            )
-                            .first()
+                        dataset_document_stmt = select(DatasetDocument).where(
+                            DatasetDocument.id == segment.document_id,
+                            DatasetDocument.enabled == True,
+                            DatasetDocument.archived == False,
                         )
+                        document = db.session.scalar(dataset_document_stmt)
                         if dataset and document:
                             source = RetrievalSourceMetadata(
                                 dataset_id=dataset.id,
@@ -327,7 +325,8 @@ class DatasetRetrieval:
 
         if dataset_id:
             # get retrieval model config
-            dataset = db.session.query(Dataset).where(Dataset.id == dataset_id).first()
+            dataset_stmt = select(Dataset).where(Dataset.id == dataset_id)
+            dataset = db.session.scalar(dataset_stmt)
             if dataset:
                 results = []
                 if dataset.provider == "external":
@@ -514,22 +513,18 @@ class DatasetRetrieval:
         dify_documents = [document for document in documents if document.provider == "dify"]
         for document in dify_documents:
             if document.metadata is not None:
-                dataset_document = (
-                    db.session.query(DatasetDocument)
-                    .where(DatasetDocument.id == document.metadata["document_id"])
-                    .first()
+                dataset_document_stmt = select(DatasetDocument).where(
+                    DatasetDocument.id == document.metadata["document_id"]
                 )
+                dataset_document = db.session.scalar(dataset_document_stmt)
                 if dataset_document:
                     if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX:
-                        child_chunk = (
-                            db.session.query(ChildChunk)
-                            .where(
-                                ChildChunk.index_node_id == document.metadata["doc_id"],
-                                ChildChunk.dataset_id == dataset_document.dataset_id,
-                                ChildChunk.document_id == dataset_document.id,
-                            )
-                            .first()
+                        child_chunk_stmt = select(ChildChunk).where(
+                            ChildChunk.index_node_id == document.metadata["doc_id"],
+                            ChildChunk.dataset_id == dataset_document.dataset_id,
+                            ChildChunk.document_id == dataset_document.id,
                         )
+                        child_chunk = db.session.scalar(child_chunk_stmt)
                         if child_chunk:
                             segment = (
                                 db.session.query(DocumentSegment)
@@ -600,7 +595,8 @@ class DatasetRetrieval:
     ):
         with flask_app.app_context():
             with Session(db.engine) as session:
-                dataset = session.query(Dataset).where(Dataset.id == dataset_id).first()
+                dataset_stmt = select(Dataset).where(Dataset.id == dataset_id)
+                dataset = db.session.scalar(dataset_stmt)
 
             if not dataset:
                 return []
@@ -647,7 +643,7 @@ class DatasetRetrieval:
                             retrieval_method=retrieval_model["search_method"],
                             dataset_id=dataset.id,
                             query=query,
-                            top_k=retrieval_model.get("top_k") or 2,
+                            top_k=retrieval_model.get("top_k") or 4,
                             score_threshold=retrieval_model.get("score_threshold", 0.0)
                             if retrieval_model["score_threshold_enabled"]
                             else 0.0,
@@ -685,7 +681,8 @@ class DatasetRetrieval:
         available_datasets = []
         for dataset_id in dataset_ids:
             # get dataset from dataset id
-            dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first()
+            dataset_stmt = select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id)
+            dataset = db.session.scalar(dataset_stmt)
 
             # pass if dataset is not available
             if not dataset:
@@ -743,7 +740,7 @@ class DatasetRetrieval:
             tool = DatasetMultiRetrieverTool.from_dataset(
                 dataset_ids=[dataset.id for dataset in available_datasets],
                 tenant_id=tenant_id,
-                top_k=retrieve_config.top_k or 2,
+                top_k=retrieve_config.top_k or 4,
                 score_threshold=retrieve_config.score_threshold,
                 hit_callbacks=[hit_callback],
                 return_resource=return_resource,
@@ -958,7 +955,8 @@ class DatasetRetrieval:
         self, dataset_ids: list, query: str, tenant_id: str, user_id: str, metadata_model_config: ModelConfig
     ) -> Optional[list[dict[str, Any]]]:
         # get all metadata field
-        metadata_fields = db.session.query(DatasetMetadata).where(DatasetMetadata.dataset_id.in_(dataset_ids)).all()
+        metadata_stmt = select(DatasetMetadata).where(DatasetMetadata.dataset_id.in_(dataset_ids))
+        metadata_fields = db.session.scalars(metadata_stmt).all()
         all_metadata_fields = [metadata_field.name for metadata_field in metadata_fields]
         # get metadata model config
         if metadata_model_config is None: