Support non-json serializable objects in LLMContext

conda/environments/all_cuda-121_arch-x86_64.yaml

@@ -63,7 +63,7 @@ dependencies:
 - networkx=2.8.8
 - newspaper3k=0.2
 - ninja=1.11
-- nlohmann_json=3.9
+- nlohmann_json=3.11
 - nodejs=18.*
 - numexpr
 - numpydoc=1.5

conda/environments/dev_cuda-121_arch-x86_64.yaml

@@ -49,7 +49,7 @@ dependencies:
 - nbsphinx
 - networkx=2.8.8
 - ninja=1.11
-- nlohmann_json=3.9
+- nlohmann_json=3.11
 - nodejs=18.*
 - numpydoc=1.5
 - nvtabular=23.08.00

dependencies.yaml

@@ -181,7 +181,7 @@ dependencies:
           - librdkafka>=1.9.2,<1.10.0a0
           - mrc=24.03
           - ninja=1.11
-          - nlohmann_json=3.9
+          - nlohmann_json=3.11
           - pkg-config=0.29 # for mrc cmake
           - protobuf=4.24
           - pybind11-stubgen=0.10.5

morpheus.code-workspace

@@ -585,6 +585,34 @@
                 },
                 "type": "cppdbg"
             },
+            {
+                "MIMode": "gdb",
+                "args": [
+                    // Uncomment to run a specific test
+                    // "--gtest_filter='TestLLMContext.PopSelectMultipleOutputs'"
+                ],
+                "cwd": "${workspaceFolder}",
+                "environment": [
+                    {
+                        "name": "GLOG_v",
+                        "value": "10"
+                    }
+                ],
+                "externalConsole": false,
+                "miDebuggerPath": "gdb",
+                "name": "Debug LLM C++ Tests",
+                "program": "${workspaceFolder}/build/morpheus/_lib/tests/test_llm.x",
+                "request": "launch",
+                "setupCommands": [
+                    {
+                        "description": "Enable pretty-printing for gdb",
+                        "ignoreFailures": true,
+                        "text": "-enable-pretty-printing"
+                    }
+                ],
+                "stopAtEntry": false,
+                "type": "cppdbg"
+            },
             {
                 "args": [
                     "--input_file=validation.csv",

morpheus/_lib/include/morpheus/llm/llm_context.hpp

@@ -18,13 +18,12 @@
 #pragma once
 
 #include "morpheus/export.h"
+#include "morpheus/llm/fwd.hpp"  // for ControlMessage
 #include "morpheus/llm/input_map.hpp"
 #include "morpheus/llm/llm_task.hpp"
-#include "morpheus/messages/control.hpp"
-#include "morpheus/utilities/json_types.hpp"
 
 #include <mrc/types.hpp>
-#include <nlohmann/json.hpp>
+#include <pymrc/utilities/json_values.hpp>
 
 #include <memory>
 #include <string>
@@ -36,7 +35,6 @@ struct LLMContextState
 {
     LLMTask task;
     std::shared_ptr<ControlMessage> message;
-    nlohmann::json values;
 };
 
 /**
@@ -110,11 +108,11 @@ class MORPHEUS_EXPORT LLMContext : public std::enable_shared_from_this<LLMContex
     std::shared_ptr<ControlMessage>& message() const;
 
     /**
-     * @brief Get all output mappings for this context.
+     * @brief Get all outputs for this context.
      *
-     * @return nlohmann::json::const_reference
+     * @return const mrc::pymrc::JSONValues&
      */
-    nlohmann::json::const_reference all_outputs() const;
+    const mrc::pymrc::JSONValues& all_outputs() const;
 
     /**
      * @brief Get full name of context containing parents up to root.
@@ -142,39 +140,39 @@ class MORPHEUS_EXPORT LLMContext : public std::enable_shared_from_this<LLMContex
     /**
      * @brief Get the input value from parent context corresponding to first internal input of this context.
      *
-     * @return nlohmann::json::const_reference
+     * @return mrc::pymrc::JSONValues
      */
-    nlohmann::json::const_reference get_input() const;
+    mrc::pymrc::JSONValues get_input() const;
 
     /**
      * @brief Get the parent output value corresponding to given internal input name.
      *
      * @param node_name internal input name
-     * @return nlohmann::json::const_reference
+     * @return mrc::pymrc::JSONValues
      */
-    nlohmann::json::const_reference get_input(const std::string& node_name) const;
+    mrc::pymrc::JSONValues get_input(const std::string& node_name) const;
 
     /**
      * @brief Get parent output values corresponding to all internal input names.
      *
-     * @return nlohmann::json
+     * @return mrc::pymrc::JSONValues
      */
-    nlohmann::json get_inputs() const;
+    mrc::pymrc::JSONValues get_inputs() const;
 
     /**
      * @brief Set output mappings for this context.
      *
      * @param outputs output mappings
      */
-    void set_output(nlohmann::json outputs);
+    void set_output(mrc::pymrc::JSONValues&& outputs);
 
     /**
      * @brief Set an output value for this context.
      *
      * @param output_name output name
      * @param output output value
      */
-    void set_output(const std::string& output_name, nlohmann::json output);
+    void set_output(const std::string& output_name, mrc::pymrc::JSONValues&& output);
 
     /**
      * @brief Set the output names to propagate from this context when using pop.
@@ -185,17 +183,24 @@ class MORPHEUS_EXPORT LLMContext : public std::enable_shared_from_this<LLMContex
 
     void outputs_complete();
 
-    nlohmann::json::const_reference view_outputs() const;
+    /**
+     * @brief Get all outputs for this context.
+     *
+     * @return const mrc::pymrc::JSONValues&
+     */
+    const mrc::pymrc::JSONValues& view_outputs() const;
 
   private:
+    input_mappings_t::const_iterator find_input(const std::string& node_name, bool throw_if_not_found = true) const;
+
     std::shared_ptr<LLMContext> m_parent{nullptr};
     std::string m_name;
     input_mappings_t m_inputs;
     std::vector<std::string> m_output_names;  // Names of keys to be used as the output. Empty means use all keys
 
     std::shared_ptr<LLMContextState> m_state;
 
-    nlohmann::json m_outputs;
+    mrc::pymrc::JSONValues m_outputs;
 
     mrc::Promise<void> m_outputs_promise;
     mrc::SharedFuture<void> m_outputs_future;

morpheus/_lib/include/morpheus/llm/llm_engine.hpp

@@ -18,12 +18,9 @@
 #pragma once
 
 #include "morpheus/export.h"
-#include "morpheus/llm/fwd.hpp"
+#include "morpheus/llm/fwd.hpp"  // for ControlMessage, LLMContext, LLMTask, LLMTaskHandler, LLMTaskHandlerRunner
 #include "morpheus/llm/input_map.hpp"
-#include "morpheus/llm/llm_context.hpp"
 #include "morpheus/llm/llm_node.hpp"
-#include "morpheus/llm/llm_task_handler.hpp"
-#include "morpheus/messages/control.hpp"
 #include "morpheus/types.hpp"
 
 #include <memory>

morpheus/_lib/include/morpheus/llm/llm_lambda_node.hpp

@@ -78,7 +78,7 @@ class LLMLambdaNode : public LLMNodeBase
         {
             const auto& arg = context->get_input();
 
-            auto output = co_await this->m_function(arg.get<std::tuple_element_t<0, args_tuple_t>>());
+            auto output = co_await this->m_function(arg.view_json().get<std::tuple_element_t<0, args_tuple_t>>());
 
             nlohmann::json outputs_json = std::move(output);
 
@@ -91,7 +91,7 @@ class LLMLambdaNode : public LLMNodeBase
         {
             auto args = context->get_inputs();
 
-            auto outputs = co_await this->m_function(args);
+            auto outputs = co_await this->m_function(args.view_json());
 
             nlohmann::json outputs_json = std::move(outputs);
 

morpheus/_lib/include/morpheus/pybind11/json.hpp

@@ -35,31 +35,36 @@ struct type_caster<nlohmann::json>
 {
   public:
     /**
-     * This macro establishes the name 'inty' in
-     * function signatures and declares a local variable
-     * 'value' of type inty
+     * This macro establishes a local variable 'value' of type nlohmann::json
      */
     PYBIND11_TYPE_CASTER(nlohmann::json, _("object"));
 
     /**
-     * Conversion part 1 (Python->C++): convert a PyObject into a inty
+     * Conversion part 1 (Python->C++): convert a PyObject into an nlohmann::json
      * instance or return false upon failure. The second argument
      * indicates whether implicit conversions should be applied.
      */
     bool load(handle src, bool convert)
     {
-        if (!src || src.is_none())
+        if (!src)
         {
             return false;
         }
 
-        value = mrc::pymrc::cast_from_pyobject(pybind11::reinterpret_borrow<pybind11::object>(src));
+        if (src.is_none())
+        {
+            value = nlohmann::json(nullptr);
+        }
+        else
+        {
+            value = mrc::pymrc::cast_from_pyobject(pybind11::reinterpret_borrow<pybind11::object>(src));
+        }
 
         return true;
     }
 
     /**
-     * Conversion part 2 (C++ -> Python): convert an inty instance into
+     * Conversion part 2 (C++ -> Python): convert an nlohmann::json instance into
      * a Python object. The second and third arguments are used to
      * indicate the return value policy and parent object (for
      * ``return_value_policy::reference_internal``) and are generally
@@ -76,14 +81,12 @@ struct type_caster<nlohmann::json_dict>
 {
   public:
     /**
-     * This macro establishes the name 'inty' in
-     * function signatures and declares a local variable
-     * 'value' of type inty
+     * This macro establishes a local variable 'value' of type nlohmann::json_dict
      */
     PYBIND11_TYPE_CASTER(nlohmann::json_dict, _("dict[str, typing.Any]"));
 
     /**
-     * Conversion part 1 (Python->C++): convert a PyObject into a inty
+     * Conversion part 1 (Python->C++): convert a PyObject into an nlohmann::json_dict
      * instance or return false upon failure. The second argument
      * indicates whether implicit conversions should be applied.
      */
@@ -106,7 +109,7 @@ struct type_caster<nlohmann::json_dict>
     }
 
     /**
-     * Conversion part 2 (C++ -> Python): convert an inty instance into
+     * Conversion part 2 (C++ -> Python): convert an nlohmann::json_dict instance into
      * a Python object. The second and third arguments are used to
      * indicate the return value policy and parent object (for
      * ``return_value_policy::reference_internal``) and are generally
@@ -123,14 +126,12 @@ struct type_caster<nlohmann::json_list>
 {
   public:
     /**
-     * This macro establishes the name 'inty' in
-     * function signatures and declares a local variable
-     * 'value' of type inty
+     * This macro establishes a local variable 'value' of type nlohmann::json_list
      */
     PYBIND11_TYPE_CASTER(nlohmann::json_list, _("list[typing.Any]"));
 
     /**
-     * Conversion part 1 (Python->C++): convert a PyObject into a inty
+     * Conversion part 1 (Python->C++): convert a PyObject into an nlohmann::json_list
      * instance or return false upon failure. The second argument
      * indicates whether implicit conversions should be applied.
      */
@@ -153,7 +154,7 @@ struct type_caster<nlohmann::json_list>
     }
 
     /**
-     * Conversion part 2 (C++ -> Python): convert an inty instance into
+     * Conversion part 2 (C++ -> Python): convert an nlohmann::json_list instance into
      * a Python object. The second and third arguments are used to
      * indicate the return value policy and parent object (for
      * ``return_value_policy::reference_internal``) and are generally

morpheus/_lib/llm/__init__.pyi

@@ -67,7 +67,7 @@ class LLMContext():
     def get_input(self) -> object: ...
     @typing.overload
     def get_input(self, node_name: str) -> object: ...
-    def get_inputs(self) -> dict: ...
+    def get_inputs(self) -> object: ...
     def message(self) -> morpheus._lib.messages.ControlMessage: ...
     def push(self, name: str, inputs: typing.List[InputMap]) -> LLMContext: ...
     @typing.overload
@@ -210,7 +210,7 @@ class LLMTaskHandler():
     def __init__(self) -> None: ...
     def get_input_names(self) -> typing.List[str]: 
         """
-        Get the input names for the task handler. 
+        Get the input names for the task handler.
 
         Returns
         -------

morpheus/_lib/llm/include/py_llm_engine.hpp

@@ -19,7 +19,7 @@
 
 #include "py_llm_node.hpp"
 
-#include "morpheus/llm/fwd.hpp"
+#include "morpheus/llm/fwd.hpp"  // for LLMTaskHandler, ControlMessage, LLMContext, LLMEngine, LLMTask
 #include "morpheus/llm/input_map.hpp"
 #include "morpheus/llm/llm_engine.hpp"
 

morpheus/_lib/llm/include/py_llm_lambda_node.hpp

@@ -17,7 +17,7 @@
 
 #pragma once
 
-#include "morpheus/llm/llm_context.hpp"
+#include "morpheus/llm/fwd.hpp"  // for LLMContext
 #include "morpheus/llm/llm_node_base.hpp"
 #include "morpheus/types.hpp"