From e90c532c3ab334274f318adfe99f6e239ee93c77 Mon Sep 17 00:00:00 2001
From: Jyong <76649700+JohnJyong@users.noreply.github.com>
Date: Thu, 17 Apr 2025 18:05:15 +0800
Subject: [PATCH] fix retrival resource miss in chatflow (#18307)

---
 api/controllers/web/message.py                |  1 +
 .../index_tool_callback_handler.py            | 24 -------------------
 api/models/model.py                           |  7 +-----
 3 files changed, 2 insertions(+), 30 deletions(-)

diff --git a/api/controllers/web/message.py b/api/controllers/web/message.py
index 494b357d4..17e9a3990 100644
--- a/api/controllers/web/message.py
+++ b/api/controllers/web/message.py
@@ -46,6 +46,7 @@ class MessageListApi(WebApiResource):
         "retriever_resources": fields.List(fields.Nested(retriever_resource_fields)),
         "created_at": TimestampField,
         "agent_thoughts": fields.List(fields.Nested(agent_thought_fields)),
+        "metadata": fields.Raw(attribute="message_metadata_dict"),
         "status": fields.String,
         "error": fields.String,
     }
diff --git a/api/core/callback_handler/index_tool_callback_handler.py b/api/core/callback_handler/index_tool_callback_handler.py
index 64c734f62..56859df7f 100644
--- a/api/core/callback_handler/index_tool_callback_handler.py
+++ b/api/core/callback_handler/index_tool_callback_handler.py
@@ -6,7 +6,6 @@ from core.rag.models.document import Document
 from extensions.ext_database import db
 from models.dataset import ChildChunk, DatasetQuery, DocumentSegment
 from models.dataset import Document as DatasetDocument
-from models.model import DatasetRetrieverResource
 
 
 class DatasetIndexToolCallbackHandler:
@@ -71,29 +70,6 @@ class DatasetIndexToolCallbackHandler:
 
     def return_retriever_resource_info(self, resource: list):
         """Handle return_retriever_resource_info."""
-        if resource and len(resource) > 0:
-            for item in resource:
-                dataset_retriever_resource = DatasetRetrieverResource(
-                    message_id=self._message_id,
-                    position=item.get("position") or 0,
-                    dataset_id=item.get("dataset_id"),
-                    dataset_name=item.get("dataset_name"),
-                    document_id=item.get("document_id"),
-                    document_name=item.get("document_name"),
-                    data_source_type=item.get("data_source_type"),
-                    segment_id=item.get("segment_id"),
-                    score=item.get("score") if "score" in item else None,
-                    hit_count=item.get("hit_count") if "hit_count" in item else None,
-                    word_count=item.get("word_count") if "word_count" in item else None,
-                    segment_position=item.get("segment_position") if "segment_position" in item else None,
-                    index_node_hash=item.get("index_node_hash") if "index_node_hash" in item else None,
-                    content=item.get("content"),
-                    retriever_from=item.get("retriever_from"),
-                    created_by=self._user_id,
-                )
-                db.session.add(dataset_retriever_resource)
-                db.session.commit()
-
         self._queue_manager.publish(
             QueueRetrieverResourcesEvent(retriever_resources=resource), PublishFrom.APPLICATION_MANAGER
         )
diff --git a/api/models/model.py b/api/models/model.py
index a826d13e7..6577492d1 100644
--- a/api/models/model.py
+++ b/api/models/model.py
@@ -1091,12 +1091,7 @@ class Message(db.Model):  # type: ignore[name-defined]
 
     @property
     def retriever_resources(self):
-        return (
-            db.session.query(DatasetRetrieverResource)
-            .filter(DatasetRetrieverResource.message_id == self.id)
-            .order_by(DatasetRetrieverResource.position.asc())
-            .all()
-        )
+        return self.message_metadata_dict.get("retriever_resources") if self.message_metadata else []
 
     @property
     def message_files(self):