Optimize knowledge retrieval performance by batching dataset quries. (#4917)

2024-06-05 13:30:32 +08:00
parent 3006124e6d
commit 7749b71fff
2 changed files with 27 additions and 19 deletions
--- a/api/core/rag/retrieval/dataset_retrieval.py
+++ b/api/core/rag/retrieval/dataset_retrieval.py
@@ -329,6 +329,7 @@ class DatasetRetrieval:
        """
        if not query:
            return
+        dataset_queries = []
        for dataset_id in dataset_ids:
            dataset_query = DatasetQuery(
                dataset_id=dataset_id,
@@ -338,7 +339,9 @@ class DatasetRetrieval:
                created_by_role=user_from,
                created_by=user_id
            )
-            db.session.add(dataset_query)
+            dataset_queries.append(dataset_query)
+        if dataset_queries:
+            db.session.add_all(dataset_queries)
        db.session.commit()

    def _retriever(self, flask_app: Flask, dataset_id: str, query: str, top_k: int, all_documents: list):