chore: DocumentSegment to Typebase (#35635)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/core/rag/datasource/keyword/jieba/jieba.py

@@ -245,6 +245,7 @@ class Jieba(BaseKeyword):
             segment = pre_segment_data["segment"]
             if pre_segment_data["keywords"]:
                 segment.keywords = pre_segment_data["keywords"]
+                assert segment.index_node_id
                 keyword_table = self._add_text_to_keyword_table(
                     keyword_table or {}, segment.index_node_id, pre_segment_data["keywords"]
                 )
@@ -253,6 +254,7 @@ class Jieba(BaseKeyword):
 
                 keywords = keyword_table_handler.extract_keywords(segment.content, keyword_number)
                 segment.keywords = list(keywords)
+                assert segment.index_node_id
                 keyword_table = self._add_text_to_keyword_table(
                     keyword_table or {}, segment.index_node_id, list(keywords)
                 )

api/core/rag/datasource/retrieval_service.py

@@ -1,5 +1,6 @@
 import concurrent.futures
 import logging
+from collections.abc import Sequence
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, NotRequired, TypedDict
 
@@ -526,7 +527,7 @@ class RetrievalService:
             index_node_ids = [i for i in index_node_ids if i]
 
             segment_ids: list[str] = []
-            index_node_segments: list[DocumentSegment] = []
+            index_node_segments: Sequence[DocumentSegment] = []
             segments: list[DocumentSegment] = []
             attachment_map: dict[str, list[AttachmentInfoDict]] = {}
             child_chunk_map: dict[str, list[ChildChunk]] = {}
@@ -568,8 +569,9 @@ class RetrievalService:
                         DocumentSegment.status == "completed",
                         DocumentSegment.index_node_id.in_(index_node_ids),
                     )
-                    index_node_segments = session.execute(document_segment_stmt).scalars().all()  # type: ignore
+                    index_node_segments = session.execute(document_segment_stmt).scalars().all()
                     for index_node_segment in index_node_segments:
+                        assert index_node_segment.index_node_id
                         doc_segment_map[index_node_segment.id] = [index_node_segment.index_node_id]
 
                 if segment_ids:

api/core/rag/docstore/dataset_docstore.py

@@ -50,6 +50,7 @@ class DatasetDocumentStore:
 
         output = {}
         for document_segment in document_segments:
+            assert document_segment.index_node_id
             doc_id = document_segment.index_node_id
             output[doc_id] = Document(
                 page_content=document_segment.content,
@@ -103,7 +104,7 @@ class DatasetDocumentStore:
 
             if not segment_document:
                 max_position += 1
-
+                assert self._document_id
                 segment_document = DocumentSegment(
                     tenant_id=self._dataset.tenant_id,
                     dataset_id=self._dataset.id,

api/core/rag/index_processor/index_processor.py

@@ -84,7 +84,7 @@ class IndexProcessor:
                     select(DocumentSegment).where(DocumentSegment.document_id == original_document_id)
                 ).all()
                 if segments:
-                    index_node_ids = [segment.index_node_id for segment in segments]
+                    index_node_ids = [segment.index_node_id for segment in segments if segment.index_node_id]
 
         indexing_start_at = time.perf_counter()
         # delete from vector index