feat: mypy for all type check (#10921)

api/services/hit_testing_service.py

@@ -1,5 +1,6 @@
 import logging
 import time
+from typing import Any
 
 from core.rag.datasource.retrieval_service import RetrievalService
 from core.rag.models.document import Document
@@ -24,7 +25,7 @@ class HitTestingService:
         dataset: Dataset,
         query: str,
         account: Account,
-        retrieval_model: dict,
+        retrieval_model: Any,  # FIXME drop this any
         external_retrieval_model: dict,
         limit: int = 10,
     ) -> dict:
@@ -68,7 +69,7 @@ class HitTestingService:
         db.session.add(dataset_query)
         db.session.commit()
 
-        return cls.compact_retrieve_response(dataset, query, all_documents)
+        return dict(cls.compact_retrieve_response(dataset, query, all_documents))
 
     @classmethod
     def external_retrieve(
@@ -102,13 +103,16 @@ class HitTestingService:
         db.session.add(dataset_query)
         db.session.commit()
 
-        return cls.compact_external_retrieve_response(dataset, query, all_documents)
+        return dict(cls.compact_external_retrieve_response(dataset, query, all_documents))
 
     @classmethod
     def compact_retrieve_response(cls, dataset: Dataset, query: str, documents: list[Document]):
         records = []
 
         for document in documents:
+            if document.metadata is None:
+                continue
+
             index_node_id = document.metadata["doc_id"]
 
             segment = (
@@ -140,7 +144,7 @@ class HitTestingService:
         }
 
     @classmethod
-    def compact_external_retrieve_response(cls, dataset: Dataset, query: str, documents: list):
+    def compact_external_retrieve_response(cls, dataset: Dataset, query: str, documents: list) -> dict[Any, Any]:
         records = []
         if dataset.provider == "external":
             for document in documents:
@@ -152,11 +156,10 @@ class HitTestingService:
                 }
                 records.append(record)
             return {
-                "query": {
-                    "content": query,
-                },
+                "query": {"content": query},
                 "records": records,
             }
+        return {"query": {"content": query}, "records": []}
 
     @classmethod
     def hit_testing_args_check(cls, args):