Underground-Digital/Workflow-Engine

1from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom2from core.app.entities.app_invoke_entities import InvokeFrom3from core.app.entities.queue_entities import QueueRetrieverResourcesEvent4from core.rag.models.document import Document5from extensions.ext_database import db6from models.dataset import DatasetQuery, DocumentSegment7from models.model import DatasetRetrieverResource8 9 10class DatasetIndexToolCallbackHandler:11    """Callback handler for dataset tool."""12 13    def __init__(14        self, queue_manager: AppQueueManager, app_id: str, message_id: str, user_id: str, invoke_from: InvokeFrom15    ) -> None:16        self._queue_manager = queue_manager17        self._app_id = app_id18        self._message_id = message_id19        self._user_id = user_id20        self._invoke_from = invoke_from21 22    def on_query(self, query: str, dataset_id: str) -> None:23        """24        Handle query.25        """26        dataset_query = DatasetQuery(27            dataset_id=dataset_id,28            content=query,29            source="app",30            source_app_id=self._app_id,31            created_by_role=(32                "account" if self._invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER} else "end_user"33            ),34            created_by=self._user_id,35        )36 37        db.session.add(dataset_query)38        db.session.commit()39 40    def on_tool_end(self, documents: list[Document]) -> None:41        """Handle tool end."""42        for document in documents:43            query = db.session.query(DocumentSegment).filter(44                DocumentSegment.index_node_id == document.metadata["doc_id"]45            )46 47            if "dataset_id" in document.metadata:48                query = query.filter(DocumentSegment.dataset_id == document.metadata["dataset_id"])49 50            # add hit count to document segment51            query.update({DocumentSegment.hit_count: DocumentSegment.hit_count + 1}, synchronize_session=False)52 53            db.session.commit()54 55    def return_retriever_resource_info(self, resource: list):56        """Handle return_retriever_resource_info."""57        if resource and len(resource) > 0:58            for item in resource:59                dataset_retriever_resource = DatasetRetrieverResource(60                    message_id=self._message_id,61                    position=item.get("position") or 0,62                    dataset_id=item.get("dataset_id"),63                    dataset_name=item.get("dataset_name"),64                    document_id=item.get("document_id"),65                    document_name=item.get("document_name"),66                    data_source_type=item.get("data_source_type"),67                    segment_id=item.get("segment_id"),68                    score=item.get("score") if "score" in item else None,69                    hit_count=item.get("hit_count") if "hit_count" in item else None,70                    word_count=item.get("word_count") if "word_count" in item else None,71                    segment_position=item.get("segment_position") if "segment_position" in item else None,72                    index_node_hash=item.get("index_node_hash") if "index_node_hash" in item else None,73                    content=item.get("content"),74                    retriever_from=item.get("retriever_from"),75                    created_by=self._user_id,76                )77                db.session.add(dataset_retriever_resource)78                db.session.commit()79 80        self._queue_manager.publish(81            QueueRetrieverResourcesEvent(retriever_resources=resource), PublishFrom.APPLICATION_MANAGER82        )83