refactor: replace bare dict with typed annotations in core rag module (#35097)

api/core/indexing_runner.py

@@ -737,7 +737,7 @@ class IndexingRunner:
     def _update_document_index_status(
         document_id: str,
         after_indexing_status: IndexingStatus,
-        extra_update_params: dict[Any, Any] | None = None,
+        extra_update_params: Mapping[Any, Any] | None = None,
     ):
         """
         Update the document indexing status.
@@ -764,7 +764,7 @@ class IndexingRunner:
         db.session.commit()
 
     @staticmethod
-    def _update_segments_by_document(dataset_document_id: str, update_params: dict[Any, Any]):
+    def _update_segments_by_document(dataset_document_id: str, update_params: Mapping[Any, Any]):
         """
         Update the document segment by document id.
         """

api/core/rag/embedding/cached_embedding.py

@@ -106,7 +106,7 @@ class CacheEmbedding(Embeddings):
 
         return text_embeddings
 
-    def embed_multimodal_documents(self, multimodel_documents: list[dict]) -> list[list[float]]:
+    def embed_multimodal_documents(self, multimodel_documents: list[dict[str, Any]]) -> list[list[float]]:
         """Embed file documents."""
         # use doc embedding cache or store if not exists
         multimodel_embeddings: list[Any] = [None for _ in range(len(multimodel_documents))]

api/core/rag/embedding/embedding_base.py

@@ -11,7 +11,7 @@ class Embeddings(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    def embed_multimodal_documents(self, multimodel_documents: list[dict]) -> list[list[float]]:
+    def embed_multimodal_documents(self, multimodel_documents: list[dict[str, Any]]) -> list[list[float]]:
         """Embed file documents."""
         raise NotImplementedError
 

api/core/rag/retrieval/output_parser/react_output.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import NamedTuple, Union
+from typing import Any, NamedTuple, Union
 
 
 @dataclass
@@ -10,7 +10,7 @@ class ReactAction:
 
     tool: str
     """The name of the Tool to execute."""
-    tool_input: Union[str, dict]
+    tool_input: Union[str, dict[str, Any]]
     """The input to pass in to the Tool."""
     log: str
     """Additional information to log about the action."""
@@ -19,7 +19,7 @@ class ReactAction:
 class ReactFinish(NamedTuple):
     """The final return value of an ReactFinish."""
 
-    return_values: dict
+    return_values: dict[str, Any]
     """Dictionary of return values."""
     log: str
     """Additional information to log about the return value"""

api/core/rag/retrieval/router/multi_dataset_react_route.py

@@ -1,5 +1,5 @@
 from collections.abc import Generator, Sequence
-from typing import Union
+from typing import Any, Union
 
 from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage
 from graphon.model_runtime.entities.message_entities import PromptMessage, PromptMessageRole, PromptMessageTool
@@ -139,7 +139,7 @@ class ReactMultiDatasetRouter:
 
     def _invoke_llm(
         self,
-        completion_param: dict,
+        completion_param: dict[str, Any],
         model_instance: ModelInstance,
         prompt_messages: list[PromptMessage],
         stop: list[str],

api/core/rag/splitter/text_splitter.py

@@ -63,7 +63,7 @@ class TextSplitter(BaseDocumentTransformer, ABC):
     def split_text(self, text: str) -> list[str]:
         """Split text into multiple components."""
 
-    def create_documents(self, texts: list[str], metadatas: list[dict] | None = None) -> list[Document]:
+    def create_documents(self, texts: list[str], metadatas: list[dict[str, Any]] | None = None) -> list[Document]:
         """Create documents from a list of texts."""
         _metadatas = metadatas or [{}] * len(texts)
         documents = []