nvidia-nat 1.2.0a20250813__py3-none-any.whl → 1.2.0rc6__py3-none-any.whl

nat/cli/commands/registry/pull.py

@@ -74,7 +74,7 @@ async def pull_artifact(registry_handler_config: RegistryHandlerBaseConfig, pack
 
 @click.group(name=__name__,
              invoke_without_command=True,
-             help=("Pull AIQ Toolkit artifacts from a remote registry "
+             help=("Pull NAT artifacts from a remote registry "
                    "by package name."))
 @click.option(
     "--config_file",
@@ -88,12 +88,12 @@ async def pull_artifact(registry_handler_config: RegistryHandlerBaseConfig, pack
     "--channel",
     type=str,
     required=True,
-    help=("The remote registry channel to use when pulling the AIQ Toolkit artifact."),
+    help=("The remote registry channel to use when pulling the NAT artifact."),
 )
 @click.argument("packages", type=str)
 def pull(channel: str, config_file: str, packages: str) -> None:
     """
-    Pull AIQ Toolkit artifacts from a remote registry channel.
+    Pull NAT artifacts from a remote registry channel.
     """
 
     from nat.settings.global_settings import GlobalSettings

nat/cli/commands/registry/registry.py

@@ -25,9 +25,7 @@ from nat.cli.commands.registry.search import search
 logger = logging.getLogger(__name__)
 
 
-@click.group(name=__name__,
-             invoke_without_command=False,
-             help="Utility to configure AIQ Toolkit remote registry channels.")
+@click.group(name=__name__, invoke_without_command=False, help="Utility to configure NAT remote registry channels.")
 def registry_command(**kwargs):
     pass
 

nat/cli/commands/registry/remove.py

@@ -48,7 +48,7 @@ async def remove_artifact(registry_handler_config: RegistryHandlerBaseConfig, pa
 
 @click.group(name=__name__,
              invoke_without_command=True,
-             help=("Remove AIQ Toolkit artifact from a remote registry by name and version."))
+             help=("Remove NAT artifact from a remote registry by name and version."))
 @click.argument("packages", type=str)
 @click.option(
     "--config_file",
@@ -62,11 +62,11 @@ async def remove_artifact(registry_handler_config: RegistryHandlerBaseConfig, pa
     "--channel",
     type=str,
     required=True,
-    help=("The remote registry channel that will remove the AIQ Toolkit artifact."),
+    help=("The remote registry channel that will remove the NAT artifact."),
 )
 def remove(channel: str, config_file: str, packages: str) -> None:
     """
-    Remove AIQ Toolkit artifacts from a remote registry.
+    Remove NAT artifacts from a remote registry.
     """
 
     from nat.settings.global_settings import GlobalSettings

nat/cli/commands/registry/search.py

@@ -62,7 +62,7 @@ async def search_artifacts(  # pylint: disable=R0917
                 registry_handler.save_search_results(search_response=search_response, save_path=save_path)
 
 
-@click.group(name=__name__, invoke_without_command=True, help="Search for AIQ Toolkit artifacts from remote registry.")
+@click.group(name=__name__, invoke_without_command=True, help="Search for NAT artifacts from remote registry.")
 @click.option(
     "--config_file",
     type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path),
@@ -75,7 +75,7 @@ async def search_artifacts(  # pylint: disable=R0917
     "--channel",
     type=str,
     required=True,
-    help=("The remote registry channel to use when pulling the AIQ Toolkit artifact."),
+    help=("The remote registry channel to use when pulling the NAT artifact."),
 )
 @click.option(
     "-o",
@@ -125,7 +125,7 @@ def search(  # pylint: disable=R0917
         n_results: int,
         output_path: str) -> None:
     """
-    Search for AIQ Toolkit artifacts with the specified configuration.
+    Search for NAT artifacts with the specified configuration.
     """
 
     from nat.settings.global_settings import GlobalSettings

nat/cli/commands/start.py

@@ -149,7 +149,7 @@ class StartCommandGroup(click.Group):
 
             # Build the command parameters
             params: list[click.Parameter] = self._build_params(registered_front_end)
-            help_msg = f"Run an AIQ Toolkit workflow using the {registered_front_end.local_name} front end."
+            help_msg = f"Run a NAT workflow using the {registered_front_end.local_name} front end."
 
             cmd = click.Command(name=registered_front_end.local_name,
                                 params=params,
@@ -178,7 +178,7 @@ class StartCommandGroup(click.Group):
         # Here we need to ensure all objects are loaded before we try to create the config object
         discover_and_register_plugins(PluginTypes.CONFIG_OBJECT)
 
-        logger.info("Starting AIQ Toolkit from config file: '%s'", config_file)
+        logger.info("Starting NAT from config file: '%s'", config_file)
 
         config_dict = load_and_override_config(config_file, override)
 
@@ -238,9 +238,9 @@ class StartCommandGroup(click.Group):
 
 @click.command(name=__name__,
                invoke_without_command=False,
-               help="Run an AIQ Toolkit workflow using a front end configuration.",
+               help="Run a NAT workflow using a front end configuration.",
                cls=StartCommandGroup)
 @click.pass_context
 def start_command(ctx: click.Context, **kwargs) -> None:
-    """Run an AIQ Toolkit workflow using a front end configuration."""
+    """Run a NAT workflow using a front end configuration."""
     pass

nat/cli/commands/workflow/templates/pyproject.toml.j2

@@ -9,14 +9,14 @@ root = "{{ rel_path_to_repo_root}}"{% else %}requires = ["setuptools >= 64"]{% e
 name = "{{ package_name }}"
 {% if editable %}dynamic = ["version"]{% else %}version = "0.1.0"{% endif %}
 dependencies = [
-  "aiqtoolkit[langchain]",
+  "nvidia-nat[langchain]",
 ]
 requires-python = ">=3.11,<3.13"
-description = "Custom AIQ Toolkit Workflow"
+description = "Custom NeMo Agent Toolkit Workflow"
 classifiers = ["Programming Language :: Python"]
 
 {% if editable %}[tool.uv.sources]
-aiqtoolkit = { path = "{{ rel_path_to_repo_root}}", editable = true }{% endif %}
+nvidia-nat = { path = "{{ rel_path_to_repo_root}}", editable = true }{% endif %}
 
-[project.entry-points.'aiq.components']
+[project.entry-points.'nat.components']
 {{ package_name }} = "{{ package_name }}.register"

nat/cli/commands/workflow/workflow_commands.py

@@ -98,7 +98,7 @@ def find_package_root(package_name: str) -> Path | None:
 
 def get_workflow_path_from_name(workflow_name: str):
     """
-    Look up the location of an installed AIQ Toolkit workflow and retrieve the root directory of the installed workflow.
+    Look up the location of an installed NAT workflow and retrieve the root directory of the installed workflow.
 
     Args:
         workflow_name: The name of the workflow.
@@ -127,13 +127,13 @@ def get_workflow_path_from_name(workflow_name: str):
     "within. Defaults to the present working directory.")
 @click.option(
     "--description",
-    default="AIQ Toolkit function template. Please update the description.",
+    default="NAT function template. Please update the description.",
     help="""A description of the component being created. Will be used to populate the docstring and will describe the
-         component when inspecting installed components using 'aiq info component'""")
+         component when inspecting installed components using 'nat info component'""")
 # pylint: disable=missing-param-doc
 def create_command(workflow_name: str, install: bool, workflow_dir: str, description: str):
     """
-    Create a new AIQ Toolkit workflow using templates.
+    Create a new NAT workflow using templates.
 
     Args:
         workflow_name (str): The name of the new workflow.
@@ -234,7 +234,7 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip
 @click.argument('workflow_name')
 def reinstall_command(workflow_name):
     """
-    Reinstall an AIQ Toolkit workflow to update dependencies and code changes.
+    Reinstall a NAT workflow to update dependencies and code changes.
 
     Args:
         workflow_name (str): The name of the workflow to reinstall.
@@ -270,7 +270,7 @@ def reinstall_command(workflow_name):
 @click.argument('workflow_name')
 def delete_command(workflow_name: str):
     """
-    Delete an AIQ Toolkit workflow and uninstall its package.
+    Delete a NAT workflow and uninstall its package.
 
     Args:
         workflow_name (str): The name of the workflow to delete.

nat/data_models/api_server.py

@@ -33,6 +33,8 @@ from pydantic_core.core_schema import ValidationInfo
 from nat.data_models.interactive import HumanPrompt
 from nat.utils.type_converter import GlobalTypeConverter
 
+FINISH_REASONS = frozenset({'stop', 'length', 'tool_calls', 'content_filter', 'function_call'})
+
 
 class Request(BaseModel):
     """
@@ -111,7 +113,7 @@ class Message(BaseModel):
 
 class ChatRequest(BaseModel):
     """
-    ChatRequest is a data model that represents a request to the AIQ Toolkit chat API.
+    ChatRequest is a data model that represents a request to the NAT chat API.
     Fully compatible with OpenAI Chat Completions API specification.
     """
 
@@ -201,7 +203,7 @@ class Choice(BaseModel):
     delta: ChoiceDelta | None = None
     finish_reason: typing.Literal['stop', 'length', 'tool_calls', 'content_filter', 'function_call'] | None = None
     index: int
-    # logprobs: AIQChoiceLogprobs | None = None
+    # logprobs: ChoiceLogprobs | None = None
 
 
 class Usage(BaseModel):
@@ -212,7 +214,7 @@ class Usage(BaseModel):
 
 class ResponseSerializable(abc.ABC):
     """
-    AIQChatResponseSerializable is an abstract class that defines the interface for serializing output for the AIQ
+    ResponseSerializable is an abstract class that defines the interface for serializing output for the NAT
     Toolkit chat streaming API.
     """
 
@@ -235,7 +237,7 @@ class ResponseBaseModelIntermediate(BaseModel, ResponseSerializable):
 
 class ChatResponse(ResponseBaseModelOutput):
     """
-    ChatResponse is a data model that represents a response from the AIQ Toolkit chat API.
+    ChatResponse is a data model that represents a response from the NAT chat API.
     Fully compatible with OpenAI Chat Completions API specification.
     """
 
@@ -283,7 +285,7 @@ class ChatResponse(ResponseBaseModelOutput):
 
 class ChatResponseChunk(ResponseBaseModelOutput):
     """
-    ChatResponseChunk is a data model that represents a response chunk from the AIQ Toolkit chat streaming API.
+    ChatResponseChunk is a data model that represents a response chunk from the NAT chat streaming API.
     Fully compatible with OpenAI Chat Completions API specification.
     """
 
@@ -347,18 +349,21 @@ class ChatResponseChunk(ResponseBaseModelOutput):
 
         delta = ChoiceDelta(content=content, role=role) if content is not None or role is not None else ChoiceDelta()
 
-        return ChatResponseChunk(id=id_,
-                                 choices=[Choice(index=0, message=None, delta=delta, finish_reason=finish_reason)],
-                                 created=created,
-                                 model=model,
-                                 object="chat.completion.chunk",
-                                 usage=usage,
-                                 system_fingerprint=system_fingerprint)
+        final_finish_reason = finish_reason if finish_reason in FINISH_REASONS else None
+
+        return ChatResponseChunk(
+            id=id_,
+            choices=[Choice(index=0, message=None, delta=delta, finish_reason=final_finish_reason)],
+            created=created,
+            model=model,
+            object="chat.completion.chunk",
+            usage=usage,
+            system_fingerprint=system_fingerprint)
 
 
 class ResponseIntermediateStep(ResponseBaseModelIntermediate):
     """
-    ResponseSerializedStep is a data model that represents a serialized step in the AIQ Toolkit chat streaming API.
+    ResponseSerializedStep is a data model that represents a serialized step in the NAT chat streaming API.
     """
 
     # Allow extra fields in the model_config to support derived models
@@ -580,7 +585,7 @@ class WebSocketSystemInteractionMessage(BaseModel):
     timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
 
 
-# ======== AIQGenerateResponse Converters ========
+# ======== GenerateResponse Converters ========
 
 
 def _generate_response_to_str(response: GenerateResponse) -> str:
@@ -606,33 +611,35 @@ def _generate_response_to_chat_response(response: GenerateResponse) -> ChatRespo
 GlobalTypeConverter.register_converter(_generate_response_to_chat_response)
 
 
-# ======== AIQChatRequest Converters ========
-def _aiq_chat_request_to_string(data: ChatRequest) -> str:
+# ======== ChatRequest Converters ========
+def _nat_chat_request_to_string(data: ChatRequest) -> str:
     if isinstance(data.messages[-1].content, str):
         return data.messages[-1].content
     return str(data.messages[-1].content)
 
 
-GlobalTypeConverter.register_converter(_aiq_chat_request_to_string)
+GlobalTypeConverter.register_converter(_nat_chat_request_to_string)
 
 
-def _string_to_aiq_chat_request(data: str) -> ChatRequest:
+def _string_to_nat_chat_request(data: str) -> ChatRequest:
     return ChatRequest.from_string(data, model="")
 
 
-GlobalTypeConverter.register_converter(_string_to_aiq_chat_request)
+GlobalTypeConverter.register_converter(_string_to_nat_chat_request)
 
 
-# ======== AIQChatResponse Converters ========
-def _aiq_chat_response_to_string(data: ChatResponse) -> str:
-    return data.choices[0].message.content or ""
+# ======== ChatResponse Converters ========
+def _nat_chat_response_to_string(data: ChatResponse) -> str:
+    if data.choices and data.choices[0].message:
+        return data.choices[0].message.content or ""
+    return ""
 
 
-GlobalTypeConverter.register_converter(_aiq_chat_response_to_string)
+GlobalTypeConverter.register_converter(_nat_chat_response_to_string)
 
 
-def _string_to_aiq_chat_response(data: str) -> ChatResponse:
-    '''Converts a string to an AIQChatResponse object'''
+def _string_to_nat_chat_response(data: str) -> ChatResponse:
+    '''Converts a string to an ChatResponse object'''
 
     # Simulate usage
     prompt_tokens = 0
@@ -644,7 +651,7 @@ def _string_to_aiq_chat_response(data: str) -> ChatResponse:
     return ChatResponse.from_string(data, usage=usage)
 
 
-GlobalTypeConverter.register_converter(_string_to_aiq_chat_response)
+GlobalTypeConverter.register_converter(_string_to_nat_chat_response)
 
 
 def _chat_response_to_chat_response_chunk(data: ChatResponse) -> ChatResponseChunk:
@@ -669,19 +676,19 @@ def _chat_response_chunk_to_string(data: ChatResponseChunk) -> str:
 GlobalTypeConverter.register_converter(_chat_response_chunk_to_string)
 
 
-def _string_to_aiq_chat_response_chunk(data: str) -> ChatResponseChunk:
-    '''Converts a string to an AIQChatResponseChunk object'''
+def _string_to_nat_chat_response_chunk(data: str) -> ChatResponseChunk:
+    '''Converts a string to an ChatResponseChunk object'''
 
     # Build and return the response
     return ChatResponseChunk.from_string(data)
 
 
-GlobalTypeConverter.register_converter(_string_to_aiq_chat_response_chunk)
+GlobalTypeConverter.register_converter(_string_to_nat_chat_response_chunk)
 
 
 # ======== AINodeMessageChunk Converters ========
-def _ai_message_chunk_to_aiq_chat_response_chunk(data) -> ChatResponseChunk:
-    '''Converts LangChain AINodeMessageChunk to AIQChatResponseChunk'''
+def _ai_message_chunk_to_nat_chat_response_chunk(data) -> ChatResponseChunk:
+    '''Converts LangChain AINodeMessageChunk to ChatResponseChunk'''
     content = ""
     if hasattr(data, 'content') and data.content is not None:
         content = str(data.content)

nat/data_models/component_ref.py

@@ -43,7 +43,7 @@ class ComponentRefNode(HashableBaseModel):
 
     Args:
         ref_name (ComponentRef): The name of the component runtime instance.
-        component_group (ComponentGroup): The component group in an AIQ Toolkit configuration object.
+        component_group (ComponentGroup): The component group in a NAT configuration object.
     """
 
     ref_name: "ComponentRef"
@@ -70,7 +70,7 @@ class ComponentRef(str, ABC):
         """Provides the component group this ComponentRef object represents.
 
         Returns:
-            ComponentGroup: A component group of the AIQ Toolkit configuration object
+            ComponentGroup: A component group of the NAT configuration object
         """
 
         pass
@@ -82,7 +82,7 @@ class ComponentRef(str, ABC):
 
 class EmbedderRef(ComponentRef):
     """
-    A reference to an embedder in an AIQ Toolkit configuration object.
+    A reference to an embedder in a NAT configuration object.
     """
 
     @property
@@ -93,7 +93,7 @@ class EmbedderRef(ComponentRef):
 
 class FunctionRef(ComponentRef):
     """
-    A reference to a function in an AIQ Toolkit configuration object.
+    A reference to a function in a NAT configuration object.
     """
 
     @property
@@ -104,7 +104,7 @@ class FunctionRef(ComponentRef):
 
 class LLMRef(ComponentRef):
     """
-    A reference to an LLM in an AIQ Toolkit configuration object.
+    A reference to an LLM in a NAT configuration object.
     """
 
     @property
@@ -115,7 +115,7 @@ class LLMRef(ComponentRef):
 
 class MemoryRef(ComponentRef):
     """
-    A reference to a memory in an AIQ Toolkit configuration object.
+    A reference to a memory in a NAT configuration object.
     """
 
     @property
@@ -126,7 +126,7 @@ class MemoryRef(ComponentRef):
 
 class ObjectStoreRef(ComponentRef):
     """
-    A reference to an object store in an AIQ toolkit configuration object.
+    A reference to an object store in a NAT configuration object.
     """
 
     @property
@@ -137,7 +137,7 @@ class ObjectStoreRef(ComponentRef):
 
 class RetrieverRef(ComponentRef):
     """
-    A reference to a retriever in an AIQ Toolkit configuration object.
+    A reference to a retriever in a NAT configuration object.
     """
 
     @property
@@ -148,7 +148,7 @@ class RetrieverRef(ComponentRef):
 
 class AuthenticationRef(ComponentRef):
     """
-    A reference to an API Authentication Provider in an AIQ Toolkit configuration object.
+    A reference to an API Authentication Provider in a NAT configuration object.
     """
 
     @property

nat/data_models/dataset_handler.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import importlib
 import json
 import typing
 from collections.abc import Callable
@@ -68,7 +69,7 @@ class EvalDatasetBaseConfig(TypedBaseModel, BaseModelRegistryTag):
     s3: EvalS3Config | None = None
 
     remote_file_path: str | None = None  # only for s3
-    file_path: Path | str = Path(".tmp/aiq/examples/default/default.json")
+    file_path: Path | str = Path(".tmp/nat/examples/default/default.json")
 
 
 class EvalDatasetJsonConfig(EvalDatasetBaseConfig, name="json"):
@@ -78,7 +79,7 @@ class EvalDatasetJsonConfig(EvalDatasetBaseConfig, name="json"):
         return pd.read_json, {}
 
 
-def read_jsonl(file_path: FilePath, **kwargs):
+def read_jsonl(file_path: FilePath):
     with open(file_path, 'r', encoding='utf-8') as f:
         data = [json.loads(line) for line in f]
     return pd.DataFrame(data)
@@ -112,12 +113,57 @@ class EvalDatasetXlsConfig(EvalDatasetBaseConfig, name="xls"):
         return pd.read_excel, {"engine": "openpyxl"}
 
 
+class EvalDatasetCustomConfig(EvalDatasetBaseConfig, name="custom"):
+    """
+    Configuration for custom dataset type that allows users to specify
+    a custom Python function to transform their dataset into EvalInput format.
+    """
+
+    function: str  # Direct import path to function, format: "module.path.function_name"
+    kwargs: dict[str, typing.Any] = {}  # Additional arguments to pass to the custom function
+
+    def parser(self) -> tuple[Callable, dict]:
+        """
+        Load and return the custom function for dataset transformation.
+
+        Returns:
+            Tuple of (custom_function, kwargs) where custom_function transforms
+            a dataset file into an EvalInput object.
+        """
+        custom_function = self._load_custom_function()
+        return custom_function, self.kwargs
+
+    def _load_custom_function(self) -> Callable:
+        """
+        Import and return the custom function using standard Python import path.
+        """
+        if not self.function:
+            raise ValueError("Function path cannot be empty")
+
+        # Split the function path to get module and function name
+        module_path, function_name = self.function.rsplit(".", 1)
+
+        # Import the module
+        module = importlib.import_module(module_path)
+
+        # Get the function from the module
+        if not hasattr(module, function_name):
+            raise AttributeError(f"Function '{function_name}' not found in module '{module_path}'")
+
+        custom_function = getattr(module, function_name)
+
+        if not callable(custom_function):
+            raise ValueError(f"'{self.function}' is not callable")
+
+        return custom_function
+
+
 # Union model with discriminator
-EvalDatasetConfig = typing.Annotated[typing.Annotated[EvalDatasetJsonConfig, Tag(EvalDatasetJsonConfig.static_type())]
-                                     | typing.Annotated[EvalDatasetCsvConfig, Tag(EvalDatasetCsvConfig.static_type())]
-                                     | typing.Annotated[EvalDatasetXlsConfig, Tag(EvalDatasetXlsConfig.static_type())]
-                                     | typing.Annotated[EvalDatasetParquetConfig,
-                                                        Tag(EvalDatasetParquetConfig.static_type())]
-                                     | typing.Annotated[EvalDatasetJsonlConfig,
-                                                        Tag(EvalDatasetJsonlConfig.static_type())],
-                                     Discriminator(TypedBaseModel.discriminator)]
+EvalDatasetConfig = typing.Annotated[
+    typing.Annotated[EvalDatasetJsonConfig, Tag(EvalDatasetJsonConfig.static_type())]
+    | typing.Annotated[EvalDatasetCsvConfig, Tag(EvalDatasetCsvConfig.static_type())]
+    | typing.Annotated[EvalDatasetXlsConfig, Tag(EvalDatasetXlsConfig.static_type())]
+    | typing.Annotated[EvalDatasetParquetConfig, Tag(EvalDatasetParquetConfig.static_type())]
+    | typing.Annotated[EvalDatasetJsonlConfig, Tag(EvalDatasetJsonlConfig.static_type())]
+    | typing.Annotated[EvalDatasetCustomConfig, Tag(EvalDatasetCustomConfig.static_type())],
+    Discriminator(TypedBaseModel.discriminator)]