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