fix processor bugs (#1048)

Author: wzh1994

lazyllm/tools/rag/parsing_service/server.py

@@ -2,13 +2,12 @@
 import threading
 import time
 import traceback
-import cloudpickle
 from datetime import datetime
 from typing import Any, Callable, Dict, Optional
 
 from lazyllm import (
     LOG, ModuleBase, ServerModule, UrlModule, FastapiApp as app,
-    LazyLLMLaunchersBase as Launcher, once_wrapper
+    LazyLLMLaunchersBase as Launcher, load_obj, once_wrapper, dump_obj
 )
 from lazyllm.thirdparty import fastapi
 
@@ -118,7 +117,7 @@ def register_algorithm(self, name: str, store: _DocumentStore, reader: Directory
                     'node_groups': node_groups,
                     'schema_extractor': schema_extractor,
                 }
-                info_pickle = cloudpickle.dumps(info_dict)
+                info_pickle = dump_obj(info_dict)
                 with self._db_manager.get_session() as session:
                     AlgoInfo = self._db_manager.get_table_orm_class('lazyllm_algorithm')
                     existing_algorithm = session.query(AlgoInfo).filter(AlgoInfo.id == name).first()
@@ -239,7 +238,7 @@ def get_algo_group_info(self, algo_id: str) -> None:
                 if algorithm is None:
                     raise fastapi.HTTPException(status_code=404, detail=f'Invalid algo_id {algo_id}')
                 info_pickle_bytes = algorithm.get('info_pickle')
-                info = cloudpickle.loads(info_pickle_bytes)
+                info = load_obj(info_pickle_bytes)
                 store: _DocumentStore = info['store']  # type: ignore
                 node_groups = info['node_groups']
 

lazyllm/tools/rag/parsing_service/worker.py

@@ -2,10 +2,9 @@
 import time
 import traceback
 import threading
-import cloudpickle
 
 from datetime import datetime
-from lazyllm import LOG, FastapiApp as app, ModuleBase, ServerModule, once_wrapper
+from lazyllm import LOG, FastapiApp as app, ModuleBase, ServerModule, once_wrapper, load_obj
 from ..utils import BaseResponse, _get_default_db_config
 from .base import (
     FINISHED_TASK_QUEUE_TABLE_INFO, WAITING_TASK_QUEUE_TABLE_INFO,
@@ -86,7 +85,7 @@ def _get_or_create_processor(self, algo_id: str) -> _Processor:
                     display_name = algorithm.display_name
                     description = algorithm.description
                     info_pickle = algorithm.info_pickle
-                    info = cloudpickle.loads(info_pickle)
+                    info = load_obj(info_pickle)
                     store = info['store']
                     reader = info['reader']
                     node_groups = info['node_groups']

lazyllm/tools/rag/store/vector/milvus_store.py

@@ -17,6 +17,16 @@
 from ...data_type import DataType
 from ...global_metadata import GlobalMetadataDesc
 
+
+def _is_empty_embedding_value(v) -> bool:
+    if v is None:
+        return True
+    if isinstance(v, (list, tuple)):
+        return len(v) == 0
+    if isinstance(v, dict):
+        return not v
+    return False
+
 MILVUS_UPSERT_BATCH_SIZE = 500
 MILVUS_PAGINATION_OFFSET = 1000
 MILVUS_INDEX_MAX_RETRY = 3
@@ -141,10 +151,31 @@ def _client_context(self):
         finally:
             self._client_pool.release(c)
 
+    def _row_has_valid_embedding(self, d: dict) -> bool:
+        '''True if row has every required embed key with a non-empty value (Milvus requires all columns same length).'''
+        emb = d.get('embedding')
+        if not emb or not isinstance(emb, dict):
+            return False
+        for k in self._embed_datatypes:
+            if _is_empty_embedding_value(emb.get(k)):
+                return False
+        return True
+
     @override
-    def upsert(self, collection_name: str, data: List[dict]) -> bool:
+    def upsert(self, collection_name: str, data: List[dict]) -> bool:  # noqa: C901
         try:
             if not data: return True
+            # Only upsert rows that have valid embedding for every key. _serialize_data omits missing/empty
+            # embedding fields, which would make pymilvus build columns with different lengths (e.g. uid 230 vs
+            # embedding___default__ 229) and raise num_rows mismatch.
+            valid_data = [d for d in data if self._row_has_valid_embedding(d)]
+            dropped = len(data) - len(valid_data)
+            if dropped:
+                LOG.warning(f'[Milvus Store - upsert] Dropping {dropped} rows with missing/empty embedding for '
+                            f'collection {collection_name}.')
+            data = valid_data
+            if not data:
+                return True
             data_embeddings = data[0].get('embedding', {})
             if not data_embeddings: return True
             with self._client_context() as client: