Fix CI

examples/llm/agents/common.py

@@ -21,7 +21,6 @@
 from langchain.llms.openai import OpenAI
 
 from morpheus.config import Config
-from morpheus.messages import ControlMessage
 from morpheus.pipeline.linear_pipeline import LinearPipeline
 from morpheus.stages.general.monitor_stage import MonitorStage
 from morpheus.stages.output.in_memory_sink_stage import InMemorySinkStage
@@ -67,8 +66,7 @@ def build_common_pipeline(config: Config, pipe: LinearPipeline, task_payload: di
     Construct the elements of the pipeline common to the simple and kafka agent pipelines.
     This method should be called after the source stage has been set.
     """
-    pipe.add_stage(
-        DeserializeStage(config, message_type=ControlMessage, task_type="llm_engine", task_payload=task_payload))
+    pipe.add_stage(DeserializeStage(config, task_type="llm_engine", task_payload=task_payload))
 
     pipe.add_stage(MonitorStage(config, description="Source rate", unit='questions'))
 

examples/llm/completion/pipeline.py

@@ -20,7 +20,6 @@
 from morpheus.config import Config
 from morpheus.config import PipelineModes
 from morpheus.io.deserializers import read_file_to_df
-from morpheus.messages import ControlMessage
 from morpheus.pipeline.linear_pipeline import LinearPipeline
 from morpheus.stages.general.monitor_stage import MonitorStage
 from morpheus.stages.input.in_memory_source_stage import InMemorySourceStage
@@ -116,8 +115,7 @@ def pipeline(num_threads: int,
 
     pipe.set_source(InMemorySourceStage(config, dataframes=[source_df], repeat=repeat_count))
 
-    pipe.add_stage(
-        DeserializeStage(config, message_type=ControlMessage, task_type="llm_engine", task_payload=completion_task))
+    pipe.add_stage(DeserializeStage(config, task_type="llm_engine", task_payload=completion_task))
 
     pipe.add_stage(MonitorStage(config, description="Source rate", unit='questions'))
 

examples/llm/rag/standalone_pipeline.py

@@ -18,7 +18,6 @@
 
 from morpheus.config import Config
 from morpheus.config import PipelineModes
-from morpheus.messages import ControlMessage
 from morpheus.pipeline.linear_pipeline import LinearPipeline
 from morpheus.stages.general.monitor_stage import MonitorStage
 from morpheus.stages.input.in_memory_source_stage import InMemorySourceStage
@@ -114,8 +113,7 @@ def standalone(num_threads,
 
     pipe.set_source(InMemorySourceStage(config, dataframes=source_dfs, repeat=repeat_count))
 
-    pipe.add_stage(
-        DeserializeStage(config, message_type=ControlMessage, task_type="llm_engine", task_payload=completion_task))
+    pipe.add_stage(DeserializeStage(config, task_type="llm_engine", task_payload=completion_task))
 
     pipe.add_stage(MonitorStage(config, description="Source rate", unit='questions'))
 

python/morpheus/morpheus/_lib/include/morpheus/messages/control.hpp

@@ -19,6 +19,7 @@
 
 #include "morpheus/export.h"                  // for MORPHEUS_EXPORT
 #include "morpheus/messages/meta.hpp"         // for MessageMeta
+#include "morpheus/types.hpp"                 // for TensorIndex
 #include "morpheus/utilities/json_types.hpp"  // for json_t
 
 #include <pybind11/pytypes.h>  // for object, dict, list

python/morpheus/morpheus/_lib/messages/module.cpp

@@ -37,15 +37,15 @@
 #include "morpheus/objects/data_table.hpp"
 #include "morpheus/objects/mutable_table_ctx_mgr.hpp"
 #include "morpheus/pybind11/json.hpp"  // IWYU pragma: keep
+#include "morpheus/types.hpp"          // for TensorIndex
 #include "morpheus/utilities/cudf_util.hpp"
 #include "morpheus/utilities/json_types.hpp"  // for json_t
 #include "morpheus/utilities/string_util.hpp"
 #include "morpheus/version.hpp"
 
 #include <glog/logging.h>  // for COMPACT_GOOGLE_LOG_INFO, LogMessage, VLOG
 #include <mrc/edge/edge_connector.hpp>
-#include <nlohmann/json.hpp>  // for basic_json
-#include <pybind11/detail/common.h>
+#include <nlohmann/json.hpp>      // for basic_json
 #include <pybind11/functional.h>  // IWYU pragma: keep
 #include <pybind11/pybind11.h>
 #include <pybind11/pytypes.h>

python/morpheus/morpheus/_lib/src/io/data_loader.cpp

@@ -18,6 +18,7 @@
 #include "morpheus/io/data_loader.hpp"
 
 #include "morpheus/messages/control.hpp"
+#include "morpheus/utilities/json_types.hpp"  // for PythonByteContainer
 
 #include <glog/logging.h>
 

python/morpheus/morpheus/_lib/tests/io/test_data_loader.cpp

@@ -22,6 +22,7 @@
 #include "morpheus/io/loaders/file.hpp"
 #include "morpheus/io/loaders/payload.hpp"
 #include "morpheus/messages/control.hpp"
+#include "morpheus/utilities/json_types.hpp"  // for PythonByteContainer
 
 #include <gtest/gtest.h>
 #include <nlohmann/json.hpp>

python/morpheus/morpheus/messages/__init__.py

@@ -31,6 +31,7 @@
 from morpheus.messages.memory.response_memory import ResponseMemoryProbs
 from morpheus.messages.message_base import MessageBase
 from morpheus.messages.message_meta import MessageMeta
+from morpheus.messages.multi_message import MultiMessage
 from morpheus.messages.message_meta import UserMessageMeta
 
 __all__ = [
@@ -42,6 +43,7 @@
     "InferenceMemoryNLP",
     "MessageBase",
     "MessageMeta",
+    "MultiMessage",
     "RawPacketMessage",
     "ResponseMemory",
     "ResponseMemoryAE",

python/morpheus/morpheus/utils/concat_df.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import typing
-
 import pandas as pd
 
 import cudf

python/morpheus_llm/morpheus_llm/_lib/src/llm/llm_engine.cpp

@@ -23,7 +23,8 @@
 #include "morpheus_llm/llm/llm_task_handler_runner.hpp"  // for LLMTaskHandlerRunner
 #include "morpheus_llm/llm/utils.hpp"
 
-#include "morpheus/messages/control.hpp"  // for ControlMessage
+#include "morpheus/messages/control.hpp"      // for ControlMessage
+#include "morpheus/utilities/json_types.hpp"  // for PythonByteContainer
 
 #include <mrc/coroutines/task.hpp>  // IWYU pragma: keep
 #include <nlohmann/json.hpp>

python/morpheus_llm/morpheus_llm/_lib/tests/llm/test_llm_task_handler_runner.cpp

@@ -24,6 +24,7 @@
 #include "morpheus/_lib/tests/test_utils/common.hpp"  // IWYU pragma: associated
 #include "morpheus/messages/control.hpp"              // for ControlMessage
 #include "morpheus/types.hpp"
+#include "morpheus/utilities/json_types.hpp"  // for PythonByteContainer
 
 #include <gtest/gtest.h>
 #include <mrc/channel/forward.hpp>

python/morpheus_llm/morpheus_llm/_lib/tests/llm/test_utils.cpp

@@ -15,6 +15,7 @@
  * limitations under the License.
  */
 
+#include "morpheus_llm/llm/fwd.hpp"  // for morpheus
 #include "morpheus_llm/llm/input_map.hpp"
 #include "morpheus_llm/llm/utils.hpp"
 

tests/benchmarks/test_bench_agents_simple_pipeline.py

@@ -30,7 +30,6 @@
 import cudf
 
 from morpheus.config import Config
-from morpheus.messages import ControlMessage
 from morpheus.pipeline.linear_pipeline import LinearPipeline
 from morpheus.stages.input.in_memory_source_stage import InMemorySourceStage
 from morpheus.stages.output.in_memory_sink_stage import InMemorySinkStage

tests/benchmarks/test_bench_completion_pipeline.py

@@ -22,7 +22,6 @@
 
 from _utils.dataset_manager import DatasetManager
 from morpheus.config import Config
-from morpheus.messages import ControlMessage
 from morpheus.pipeline.linear_pipeline import LinearPipeline
 from morpheus.stages.input.in_memory_source_stage import InMemorySourceStage
 from morpheus.stages.output.in_memory_sink_stage import InMemorySinkStage

tests/benchmarks/test_bench_rag_standalone_pipeline.py

@@ -28,7 +28,6 @@
 from _utils.milvus import populate_milvus
 from morpheus.config import Config
 from morpheus.config import PipelineModes
-from morpheus.messages import ControlMessage
 from morpheus.pipeline.linear_pipeline import LinearPipeline
 from morpheus.stages.input.in_memory_source_stage import InMemorySourceStage
 from morpheus.stages.output.in_memory_sink_stage import InMemorySinkStage

tests/examples/ransomware_detection/test_create_features.py

@@ -24,7 +24,6 @@
 from _utils import TEST_DIRS
 from _utils.dataset_manager import DatasetManager
 from morpheus.config import Config
-from morpheus.messages import MultiMessage
 from morpheus.messages.message_meta import AppShieldMessageMeta
 from morpheus.pipeline.control_message_stage import ControlMessageStage
 from morpheus.stages.input.appshield_source_stage import AppShieldSourceStage
@@ -174,7 +173,7 @@ def test_create_multi_messages(self,
 
         prev_loc = 0
         for (i, _multi_message) in enumerate(multi_messages):
-            assert isinstance(_multi_message, MultiMessage)
+            assert isinstance(_multi_message)
             pid = pids[i]
             (_multi_message.get_meta(['pid_process']) == pid).all()
             assert _multi_message.mess_offset == prev_loc

tests/pipeline/test_pipeline.py

@@ -27,7 +27,6 @@
 from morpheus.config import Config
 from morpheus.messages import ControlMessage
 from morpheus.messages import MessageMeta
-from morpheus.messages import MultiMessage
 from morpheus.pipeline import LinearPipeline
 from morpheus.pipeline import Pipeline
 from morpheus.stages.boundary.linear_boundary_stage import LinearBoundaryEgressStage
@@ -226,7 +225,7 @@ def test_add_edge_input_port_errors(config: Config, num_inputs: int):
         pipe.add_edge(start_stage.output_ports[0], end_stage)
 
 
-@pytest.mark.parametrize("data_type", [int, float, str, MessageMeta, ControlMessage, MultiMessage])
+@pytest.mark.parametrize("data_type", [int, float, str, MessageMeta, ControlMessage])
 def test_add_segment_edge(config: Config, data_type: type):
     pipe = Pipeline(config)
 

tests/stages/test_generate_viz_frames_stage.py

@@ -21,7 +21,6 @@
 
 import cudf
 
-import morpheus._lib.messages as _messages
 from morpheus.config import Config
 from morpheus.messages import ControlMessage
 from morpheus.messages import MessageMeta

tests/stages/test_ml_flow_drift_stage.py

@@ -21,7 +21,6 @@
 import pytest
 import typing_utils
 
-import morpheus._lib.messages as _messages
 from morpheus.messages import ControlMessage
 from morpheus.messages.message_meta import MessageMeta
 from morpheus.stages.postprocess.ml_flow_drift_stage import MLFlowDriftStage

tests/test_add_scores_stage.py

@@ -22,10 +22,10 @@
 
 import cudf
 
-import morpheus._lib.messages as _messages
 from _utils.dataset_manager import DatasetManager
 from morpheus.config import Config
 from morpheus.messages import ControlMessage
+from morpheus.messages import TensorMemory
 from morpheus.messages.message_meta import MessageMeta
 from morpheus.stages.postprocess.add_classifications_stage import AddClassificationsStage
 from morpheus.stages.postprocess.add_scores_stage import AddScoresStage