diff --git a/api/core/rag/index_processor/index_processor.py b/api/core/rag/index_processor/index_processor.py index 825ae01226..813a84cbbd 100644 --- a/api/core/rag/index_processor/index_processor.py +++ b/api/core/rag/index_processor/index_processor.py @@ -12,7 +12,7 @@ from core.db.session_factory import session_factory from core.rag.index_processor.constant.index_type import IndexTechniqueType from core.rag.index_processor.index_processor_base import SummaryIndexSettingDict from core.workflow.nodes.knowledge_index.exc import KnowledgeIndexNodeError -from core.workflow.nodes.knowledge_index.protocols import Preview, PreviewItem, QaPreview +from core.workflow.nodes.knowledge_index.protocols import IndexingResultDict, Preview, PreviewItem, QaPreview from models.dataset import Dataset, Document, DocumentSegment from .index_processor_factory import IndexProcessorFactory @@ -61,7 +61,7 @@ class IndexProcessor: chunks: Mapping[str, Any], batch: Any, summary_index_setting: SummaryIndexSettingDict | None = None, - ): + ) -> IndexingResultDict: with session_factory.create_session() as session: document = session.query(Document).filter_by(id=document_id).first() if not document: @@ -129,7 +129,7 @@ class IndexProcessor: } ) - return { + result: IndexingResultDict = { "dataset_id": dataset_id, "dataset_name": dataset_name_value, "batch": batch, @@ -138,6 +138,7 @@ class IndexProcessor: "created_at": created_at_value.timestamp(), "display_status": "completed", } + return result def get_preview_output( self, diff --git a/api/core/workflow/nodes/knowledge_index/protocols.py b/api/core/workflow/nodes/knowledge_index/protocols.py index bb52123082..6668f0c98e 100644 --- a/api/core/workflow/nodes/knowledge_index/protocols.py +++ b/api/core/workflow/nodes/knowledge_index/protocols.py @@ -1,9 +1,19 @@ from collections.abc import Mapping -from typing import Any, Protocol +from typing import Any, Protocol, TypedDict from pydantic import BaseModel, Field +class IndexingResultDict(TypedDict): + dataset_id: str + dataset_name: str + batch: Any + document_id: str + document_name: str + created_at: float + display_status: str + + class PreviewItem(BaseModel): content: str | None = Field(default=None) child_chunks: list[str] | None = Field(default=None) @@ -34,7 +44,7 @@ class IndexProcessorProtocol(Protocol): chunks: Mapping[str, Any], batch: Any, summary_index_setting: dict | None = None, - ) -> dict[str, Any]: ... + ) -> IndexingResultDict: ... def get_preview_output( self, chunks: Any, dataset_id: str, document_id: str, chunk_structure: str, summary_index_setting: dict | None