Merge pull request #1321 from dagardner-nv/david-fea-sherlock-llm-gen…

…erate-node

Add docstrings & tests for LLMGenerateNode

examples/llm/vdb_upload/README.md

@@ -125,10 +125,7 @@ Before running the pipeline, we need to ensure that the following services are r
 
 - From the Morpheus repo root directory, run the following to launch Triton and load the `all-MiniLM-L6-v2` model:
   ```bash
-  docker run --rm -ti --gpus=all -p8000:8000 -p8001:8001 -p8002:8002
-   -v $PWD/models:/models nvcr.io/nvidia/tritonserver:23.06-py3 tritonserver 
-   --model-repository=/models/triton-model-repo --exit-on-error=false --model-control-mode=explicit
-   --load-model all-MiniLM-L6-v2
+  docker run --rm -ti --gpus=all -p8000:8000 -p8001:8001 -p8002:8002 -v $PWD/models:/models nvcr.io/nvidia/tritonserver:23.06-py3 tritonserver --model-repository=/models/triton-model-repo --exit-on-error=false --model-control-mode=explicit --load-model all-MiniLM-L6-v2
   ```
 
   This will launch Triton and only load the `all-MiniLM-L6-v2` model. Once Triton has loaded the model, the following

morpheus/_lib/llm/__init__.pyi

@@ -101,8 +101,29 @@ class LLMContext():
     pass
 class LLMNodeBase():
     def __init__(self) -> None: ...
-    def execute(self, context: LLMContext) -> typing.Awaitable[LLMContext]: ...
-    def get_input_names(self) -> typing.List[str]: ...
+    def execute(self, context: LLMContext) -> typing.Awaitable[LLMContext]: 
+        """
+        Execute the current node with the given `context` instance.
+
+        All inputs for the given node should be fetched from the context, typically by calling either
+        `context.get_inputs` to fetch all inputs as a `dict`, or `context.get_input` to fetch a specific input.
+
+        Similarly the output of the node is written to the context using `context.set_output`.
+
+        Parameters
+        ----------
+        context : `morpheus._lib.llm.LLMContext`
+            Context instance to use for the execution
+        """
+    def get_input_names(self) -> typing.List[str]: 
+        """
+        Get the input names for the node.
+
+        Returns
+        -------
+        list[str]
+            The input names for the node
+        """
     pass
 class LLMEngineStage(mrc.core.segment.SegmentObject):
     def __init__(self, builder: mrc.core.segment.Builder, name: str, engine: LLMEngine) -> None: ...

morpheus/_lib/llm/module.cpp

@@ -212,8 +212,33 @@ PYBIND11_MODULE(llm, _module)
 
     py::class_<LLMNodeBase, PyLLMNodeBase<>, std::shared_ptr<LLMNodeBase>>(_module, "LLMNodeBase")
         .def(py::init_alias<>())
-        .def("get_input_names", &LLMNodeBase::get_input_names)
-        .def("execute", &LLMNodeBase::execute, py::arg("context"));
+        .def("get_input_names",
+             &LLMNodeBase::get_input_names,
+             R"pbdoc(
+                Get the input names for the node.
+
+                Returns
+                -------
+                list[str]
+                    The input names for the node
+             )pbdoc")
+        .def("execute",
+             &LLMNodeBase::execute,
+             py::arg("context"),
+             R"pbdoc(
+                Execute the current node with the given `context` instance.
+
+                All inputs for the given node should be fetched from the context, typically by calling either
+                `context.get_inputs` to fetch all inputs as a `dict`, or `context.get_input` to fetch a specific input.
+
+                Similarly the output of the node is written to the context using `context.set_output`.
+
+                Parameters
+                ----------
+                context : `morpheus._lib.llm.LLMContext`
+                    Context instance to use for the execution
+
+            )pbdoc");
 
     py::class_<LLMNodeRunner, std::shared_ptr<LLMNodeRunner>>(_module, "LLMNodeRunner")
         .def_property_readonly("inputs", &LLMNodeRunner::inputs)

morpheus/llm/nodes/llm_generate_node.py

@@ -23,6 +23,15 @@
 
 
 class LLMGenerateNode(LLMNodeBase):
+    """
+    Generates responses from an LLM using the provided `llm_client` instance based on prompts provided as input from
+    upstream nodes.
+
+    Parameters
+    ----------
+    llm_client : LLMClient
+        The client instance to use to generate responses.
+    """
 
     def __init__(self, llm_client: LLMClient) -> None:
         super().__init__()

tests/llm/nodes/conftest.py

@@ -0,0 +1,27 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest import mock
+
+import pytest
+
+
+@pytest.fixture(name="mock_llm_client")
+def mock_llm_client_fixture():
+    from morpheus.llm.services.llm_service import LLMClient
+    mock_client = mock.MagicMock(LLMClient)
+    mock_client.return_value = mock_client
+    mock_client.generate_batch_async = mock.AsyncMock()
+    return mock_client

tests/llm/nodes/test_llm_generate_node.py

@@ -0,0 +1,39 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest import mock
+
+from _utils.llm import execute_node
+from morpheus.llm import LLMNodeBase
+from morpheus.llm.nodes.llm_generate_node import LLMGenerateNode
+
+
+def test_constructor(mock_llm_client: mock.MagicMock):
+    node = LLMGenerateNode(llm_client=mock_llm_client)
+    assert isinstance(node, LLMNodeBase)
+
+
+def test_get_input_names(mock_llm_client: mock.MagicMock):
+    node = LLMGenerateNode(llm_client=mock_llm_client)
+    assert node.get_input_names() == ["prompt"]
+
+
+def test_execute(mock_llm_client: mock.MagicMock):
+    expected_output = ["response1", "response2"]
+    mock_llm_client.generate_batch_async.return_value = expected_output.copy()
+
+    node = LLMGenerateNode(llm_client=mock_llm_client)
+    assert execute_node(node, prompt=["prompt1", "prompt2"]) == expected_output
+    mock_llm_client.generate_batch_async.assert_called_once_with(["prompt1", "prompt2"])

tests/llm/nodes/test_llm_generate_node_pipe.py

@@ -0,0 +1,66 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest import mock
+
+import pytest
+
+import cudf
+
+from _utils import assert_results
+from morpheus.config import Config
+from morpheus.llm import LLMEngine
+from morpheus.llm.llm_engine_stage import LLMEngineStage
+from morpheus.llm.nodes.extracter_node import ExtracterNode
+from morpheus.llm.nodes.llm_generate_node import LLMGenerateNode
+from morpheus.llm.task_handlers.simple_task_handler import SimpleTaskHandler
+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.compare_dataframe_stage import CompareDataFrameStage
+from morpheus.stages.preprocess.deserialize_stage import DeserializeStage
+
+
+def _build_engine(mock_llm_client: mock.MagicMock) -> LLMEngine:
+    engine = LLMEngine()
+    engine.add_node("extracter", node=ExtracterNode())
+    engine.add_node("generate", inputs=["/extracter"], node=LLMGenerateNode(llm_client=mock_llm_client))
+    engine.add_task_handler(inputs=["/generate"], handler=SimpleTaskHandler())
+
+    return engine
+
+
+@pytest.mark.use_python
+def test_pipeline(config: Config, mock_llm_client: mock.MagicMock):
+    expected_output = ["response1", "response2"]
+    mock_llm_client.generate_batch_async.return_value = expected_output.copy()
+
+    values = {'prompt': ["prompt1", "prompt2"]}
+    input_df = cudf.DataFrame(values)
+    expected_df = input_df.copy(deep=True)
+    expected_df["response"] = expected_output
+
+    task_payload = {"task_type": "llm_engine", "task_dict": {"input_keys": sorted(values.keys())}}
+
+    pipe = LinearPipeline(config)
+    pipe.set_source(InMemorySourceStage(config, dataframes=[input_df]))
+    pipe.add_stage(
+        DeserializeStage(config, message_type=ControlMessage, task_type="llm_engine", task_payload=task_payload))
+    pipe.add_stage(LLMEngineStage(config, engine=_build_engine(mock_llm_client=mock_llm_client)))
+    sink = pipe.add_stage(CompareDataFrameStage(config, compare_df=expected_df))
+
+    pipe.run()
+
+    assert_results(sink.get_results())

tests/llm/nodes/test_prompt_template_node.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

tests/llm/nodes/test_prompt_template_node_pipe.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");