Merge remote-tracking branch 'upstream/fea-agent-morpheus' into opena…

…i-client-updates

morpheus/llm/services/nvfoundation_llm_service.py

@@ -60,10 +60,21 @@ def __init__(self, parent: "NVFoundationLLMService", *, model_name: str, **model
         self._model_kwargs = model_kwargs
         self._prompt_key = "prompt"
 
-        self._client = ChatNVIDIA(api_key=self._parent._api_key,
-                                  base_url=self._parent._base_url,
-                                  model=model_name,
-                                  **model_kwargs)  # type: ignore
+        chat_kwargs = {
+            "model": model_name,
+            "api_key": self._parent._api_key,
+            "base_url": self._parent._base_url,
+        }
+
+        # Remove None values set by the environment in the kwargs
+        if (chat_kwargs["api_key"] is None):
+            del chat_kwargs["api_key"]
+
+        if (chat_kwargs["base_url"] is None):
+            del chat_kwargs["base_url"]
+
+        # Combine the chat args with the model
+        self._client = ChatNVIDIA(**{**chat_kwargs, **model_kwargs})  # type: ignore
 
     def get_input_names(self) -> list[str]:
         schema = self._client.get_input_schema()
@@ -78,7 +89,12 @@ def generate(self, **input_dict) -> str:
         input_dict : dict
             Input containing prompt data.
         """
-        return self.generate_batch({self._prompt_key: [input_dict[self._prompt_key]]})[0]
+
+        inputs = {self._prompt_key: [input_dict[self._prompt_key]]}
+
+        input_dict.pop(self._prompt_key)
+
+        return self.generate_batch(inputs=inputs, **input_dict)[0]
 
     async def generate_async(self, **input_dict) -> str:
         """
@@ -105,7 +121,9 @@ def generate_batch(self, inputs: dict[str, list], **kwargs) -> list[str]:
         """
         prompts = [StringPromptValue(text=p) for p in inputs[self._prompt_key]]
 
-        responses = self._client.generate_prompt(prompts=prompts, **self._model_kwargs)  # type: ignore
+        final_kwargs = {**self._model_kwargs, **kwargs}
+
+        responses = self._client.generate_prompt(prompts=prompts, **final_kwargs)  # type: ignore
 
         return [g[0].text for g in responses.generations]
 

morpheus/llm/services/openai_chat_service.py

@@ -37,6 +37,8 @@
     import openai
     import openai.types.chat
     import openai.types.chat.chat_completion
+    from openai.types.chat.completion_create_params import ResponseFormat
+
 except ImportError as import_exc:
     IMPORT_EXCEPTION = import_exc
 
@@ -75,11 +77,15 @@ class OpenAIChatClient(LLMClient):
     model_name : str
         The name of the model to interact with.
 
-    set_assistant: bool, optional default=False
-        When `True`, a second input field named `assistant` will be used to proide additional context to the model.
+    set_assistant: bool, optional
+        When `True`, a second input field named `assistant` will be used to proide additional context to the model, by
+        default False
+
+    max_retries: int, optional
+        The maximum number of retries to attempt when making a request to the OpenAI API, by default 10
 
-    max_retries: int, optional default=10
-        The maximum number of retries to attempt when making a request to the OpenAI API.
+    json: bool, optional
+        When `True`, the response will be returned as a JSON object, by default False
 
     model_kwargs : dict[str, typing.Any]
         Additional keyword arguments to pass to the model when generating text.
@@ -94,6 +100,7 @@ def __init__(self,
                  model_name: str,
                  set_assistant: bool = False,
                  max_retries: int = 10,
+                 json=False,
                  **model_kwargs) -> None:
         if IMPORT_EXCEPTION is not None:
             raise ImportError(IMPORT_ERROR_MESSAGE) from IMPORT_EXCEPTION
@@ -108,10 +115,14 @@ def __init__(self,
         self._set_assistant = set_assistant
         self._prompt_key = "prompt"
         self._assistant_key = "assistant"
+        self._json = json
 
         # Preserve original configuration.
         self._model_kwargs = copy.deepcopy(model_kwargs)
 
+        if (self._json):
+            self._model_kwargs["response_format"] = ResponseFormat(type="json_object")
+
         # Create the client objects for both sync and async
         self._client = openai.OpenAI(api_key=parent._api_key, base_url=parent._base_url, max_retries=max_retries)
         self._client_async = openai.AsyncOpenAI(api_key=parent._api_key,
@@ -382,6 +393,7 @@ def get_client(self,
                    model_name: str,
                    set_assistant: bool = False,
                    max_retries: int = 10,
+                   json=False,
                    **model_kwargs) -> OpenAIChatClient:
         """
         Returns a client for interacting with a specific model. This method is the preferred way to create a client.
@@ -391,11 +403,15 @@ def get_client(self,
         model_name : str
             The name of the model to create a client for.
 
-        set_assistant: bool, optional default=False
-            When `True`, a second input field named `assistant` will be used to proide additional context to the model.
+        set_assistant: bool, optional
+            When `True`, a second input field named `assistant` will be used to proide additional context to the model,
+            by default False
+
+        max_retries: int, optional
+            The maximum number of retries to attempt when making a request to the OpenAI API, by default 10
 
-        max_retries: int, optional default=10
-            The maximum number of retries to attempt when making a request to the OpenAI API.
+        json: bool, optional
+            When `True`, the response will be returned as a JSON object, by default False
 
         model_kwargs : dict[str, typing.Any]
             Additional keyword arguments to pass to the model when generating text. Arguments specified here will
@@ -408,4 +424,5 @@ def get_client(self,
                                 model_name=model_name,
                                 set_assistant=set_assistant,
                                 max_retries=max_retries,
+                                json=json,
                                 **final_model_kwargs)

morpheus/messages/multi_ae_message.py

@@ -14,10 +14,13 @@
 
 import dataclasses
 import logging
+import typing
 
 from morpheus.messages.message_meta import MessageMeta
 from morpheus.messages.multi_message import MultiMessage
-from morpheus.models.dfencoder import AutoEncoder
+
+if (typing.TYPE_CHECKING):
+    from morpheus.models.dfencoder import AutoEncoder
 
 logger = logging.getLogger(__name__)
 
@@ -28,7 +31,7 @@ class MultiAEMessage(MultiMessage):
     Subclass of `MultiMessage` specific to the AutoEncoder pipeline, which contains the model.
     """
 
-    model: AutoEncoder
+    model: "AutoEncoder"
     train_scores_mean: float
     train_scores_std: float
 
@@ -37,7 +40,7 @@ def __init__(self,
                  meta: MessageMeta,
                  mess_offset: int = 0,
                  mess_count: int = -1,
-                 model: AutoEncoder,
+                 model: "AutoEncoder",
                  train_scores_mean: float = 0.0,
                  train_scores_std: float = 1.0):
         super().__init__(meta=meta, mess_offset=mess_offset, mess_count=mess_count)

morpheus/stages/output/write_to_file_stage.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Write to file stage."""
 
+import os
 import typing
 
 import mrc
@@ -92,6 +93,9 @@ def supports_cpp_node(self):
     def _build_single(self, builder: mrc.Builder, input_node: mrc.SegmentObject) -> mrc.SegmentObject:
         # Sink to file
         if (self._build_cpp_node()):
+
+            os.makedirs(os.path.dirname(self._controller.output_file), exist_ok=True)
+
             to_file_node = _stages.WriteToFileStage(builder,
                                                     self.unique_name,
                                                     self._controller.output_file,

morpheus/utils/column_info.py

@@ -21,15 +21,16 @@
 from datetime import datetime
 from functools import partial
 
-with warnings.catch_warnings():
-    # Ignore warning regarding tensorflow not being installed
-    warnings.filterwarnings("ignore", message=".*No module named 'tensorflow'", category=UserWarning)
-    import nvtabular as nvt
-
 import pandas as pd
 
 import cudf
 
+if (typing.TYPE_CHECKING):
+    with warnings.catch_warnings():
+        # Ignore warning regarding tensorflow not being installed
+        warnings.filterwarnings("ignore", message=".*No module named 'tensorflow'", category=UserWarning)
+        import nvtabular as nvt
+
 logger = logging.getLogger(f"morpheus.{__name__}")
 
 DEFAULT_DATE = '1970-01-01T00:00:00.000000+00:00'
@@ -748,7 +749,7 @@ class DataFrameInputSchema:
     input_columns: typing.Dict[str, str] = dataclasses.field(init=False, repr=False)
     output_columns: typing.List[tuple[str, str]] = dataclasses.field(init=False, repr=False)
 
-    nvt_workflow: nvt.Workflow = dataclasses.field(init=False, repr=False)
+    nvt_workflow: "nvt.Workflow" = dataclasses.field(init=False, repr=False)
     prep_dataframe: typing.Callable[[pd.DataFrame], typing.List[str]] = dataclasses.field(init=False, repr=False)
 
     def __post_init__(self):