|
|
|
|
@ -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
|
|
|
|
|
)
|
|
|
|
|
|